u4/tests/U4SimtraceSimpleTest.py

0001 #!/usr/bin/env python
0002 """
0003 U4SimtraceSimpleTest.py
0004 =========================
0005
0006 envvars
0007 ---------
0008
0009 FOLD, SFOLD
0010     mandatory first geometry
0011 TFOLD
0012     optional second geometry
0013 AFOLD
0014     optional first simulation photon histories
0015 BFOLD
0016     optional second simulation photon histories
0017
0018 "sibling" classes
0019 --------------------
0020
0021 U4SimtraceTest
0022    intersection geometry
0023
0024 SEvt
0025    photon histories
0026
0027 """
0028 import os, logging, numpy as np
0029 log = logging.getLogger(__name__)
0030
0031 from collections import OrderedDict as odict
0032 from opticks.ana.fold import Fold
0033 from opticks.ana.pvplt import *
0034 from opticks.ana.p import *
0035 from opticks.ana.eget import eintarray_, efloatarray_
0036 from opticks.sysrap.sevt import SEvt
0037
0038
0039 COLORS = "red green blue cyan magenta yellow pink orange purple lightgreen".split()
0040 GCOL = "grey"
0041
0042 TIGHT = int(os.environ.get("TIGHT", "0")) == 1
0043 SZ = float(os.environ.get("SZ",3))
0044 REVERSE = int(os.environ.get("REVERSE","0")) == 1
0045 size = np.array([1280, 720])
0046 X,Y,Z = 0,1,2
0047
0048
0049 # use eg FOCUS=0,0,100 to select part of view with aspect ratio preserved
0050 #XLIM=efloatarray_("XLIM","0,0")
0051 #YLIM=efloatarray_("YLIM","0,0")
0052
0053 FOLD = os.environ.get("FOLD", None)
0054 SFOLD = os.environ.get("SFOLD", None)
0055 TFOLD = os.environ.get("TFOLD", None)
0056 AFOLD = os.environ.get("AFOLD", None)
0057 BFOLD = os.environ.get("BFOLD", None)
0058
0059 N = int(os.environ.get("VERSION","-1"))
0060 GEOM = os.environ.get("GEOM", "DummyGEOM")
0061 GEOMList = os.environ.get("%s_GEOMList" % GEOM, "DummyGEOMList")
0062 FOCUS = os.environ.get("FOCUS", "")
0063
0064 APID = int(os.environ.get("APID", -1 ))
0065 BPID = int(os.environ.get("BPID", -1 ))
0066 AOPT = os.environ.get("AOPT", "")
0067 BOPT = os.environ.get("BOPT", "")
0068
0069
0070
0071 topline = "N=%d " % N
0072
0073 if len(FOCUS) > 0:
0074    topline += " FOCUS=%s " % FOCUS
0075 pass
0076
0077
0078 if APID > -1: topline += "APID=%d " % APID
0079 if len(AOPT) > 0: topline += "AOPT=%s " % AOPT
0080
0081 if BPID > -1: topline += "BPID=%d " % BPID
0082 if len(BOPT) > 0: topline += "BOPT=%s " % BOPT
0083
0084 topline += " ./U4SimtraceSimpleTest.sh ana   # %s/%s " % (GEOM, GEOMList)
0085
0086 TOPLINE = os.environ.get("TOPLINE",topline )
0087 BOTLINE = os.environ.get("BOTLINE","%s" % SFOLD )
0088 THIRDLINE = os.environ.get("THIRDLINE", "APID:%d BPID:%d " % (APID,BPID) )
0089
0090 AXES = np.array(list(map(int,os.environ.get("AXES","0,2").split(","))), dtype=np.int8 )
0091 H,V = AXES
0092
0093 log = logging.getLogger(__name__)
0094 np.set_printoptions(suppress=True, edgeitems=5, linewidth=200,precision=3)
0095
0096
0097 class U4SimtraceSimpleTest(object):
0098     """
0099     MAYBE: come up with a better name for this and relocate into ana or SysRap for reuse
0100     ACTUALLY this is using trs names, so its not as general as first thought ...
0101     it is limited to small test geometries.
0102
0103     The more general GPU simtrace deserves more attention than this one.
0104
0105     HMM: there are developments elsewhere that are close to this, MAYBE: consolidate
0106     """
0107     @classmethod
0108     def Load(cls, fold, **kwa):
0109         if fold is None: return None
0110         if int(kwa.get("NEVT",0)) > 0:
0111             f = Fold.LoadConcat(fold,**kwa)
0112         else:
0113             f = Fold.Load(fold, **kwa )
0114         pass
0115         return None if f is None else cls(f)
0116
0117
0118     def __repr__(self):
0119         return "U4SimtraceSimpleTest %s:%s " % (self.symbol, self.fold.base)
0120
0121     def __init__(self, fold):
0122         """
0123         Note the container fold Loads subfold for each of the trs_names
0124         """
0125         trs_names = np.loadtxt( os.path.join(fold.base, "trs_names.txt"), dtype=np.object )
0126
0127         sfs = odict()
0128         for i, name in enumerate(trs_names):
0129             sfs[name] = Fold.Load(fold.base, name,"p001",  symbol="%s%0.2d" % (fold.symbol,i) )
0130         pass  # Fold for each solid
0131
0132         symbol = fold.symbol
0133
0134         kvol = "%sVOL" % symbol.upper()
0135         vol = eintarray_(kvol)
0136
0137         self.fold = fold
0138         self.trs_names = trs_names
0139         self.sfs = sfs
0140         self.opp = []
0141         self.kvol = kvol
0142         self.vol = vol
0143         self.symbol = symbol
0144
0145     def geoplot(self, pl):
0146         """
0147         :param pl: either tuple of (fig,ax) for MODE 2 or pyvista pl for MODE 3
0148         """
0149         fig, ax = None, None
0150         if MODE == 2:
0151             fig, axs = pl
0152             assert len(axs) == 1
0153             ax = axs[0]
0154         elif MODE == 2:
0155             pass
0156         pass
0157
0158         trs = self.fold.trs
0159         trs_names = self.trs_names
0160         sfs = self.sfs
0161         vol = self.vol
0162         kvol = self.kvol
0163
0164         vsel = len(vol) > 0   # selecting volumes by index from eg SVOL envvar
0165         num = len(trs_names)
0166
0167         if vsel:
0168             print("mp_geoplot : volume selection %s vol %s num %d " % (kvol, str(vol), num ))
0169         pass
0170
0171         for j in range(num):
0172             i = num - 1 - j  if REVERSE else j
0173
0174             soname = trs_names[i] ## solid name for every transform, lots of repeated solid names
0175
0176             if vsel:
0177                 if i in vol:
0178                     print(" %s PROCEED i %d soname %s " % (kvol, i, soname))
0179                 else:
0180                     print(" %s SKIP    i %d soname %s " % (kvol, i, soname))
0181                     continue
0182                 pass
0183             pass
0184
0185             sf = sfs[soname]
0186
0187             if not getattr(sf, 'simtrace', None) is None:
0188                 _lpos = sf.simtrace[:,1].copy()
0189                 _lpos[:,3] = 1
0190             else:
0191                  print("missing simtrace for soname:%s " % soname)
0192                  continue
0193             pass
0194
0195             gspo = sf.genstep[:,5,:3]
0196             if "GENSTEP" in os.environ:
0197                 if MODE == 2:
0198                     ax.scatter( gspo[:,H], gspo[:,V], s=SZ, color="black" )
0199                 elif MODE == 3:
0200                     pl.add_points( gspo[:,:3], color=color, label=label)
0201                 pass
0202             pass
0203
0204
0205             tr = np.float32(trs[i])
0206             tr[:,3] = [0,0,0,1]   ## fixup 4th column, as may contain identity info
0207
0208             dori = np.sqrt(np.sum(_lpos[:,:3]*_lpos[:,:3],axis=1))
0209             lpos = _lpos[dori>0]    # exclude points at local origin, misses
0210
0211             gpos = np.dot( lpos, tr )
0212             pass
0213             color = COLORS[ i % len(COLORS)]
0214
0215             label = str(soname)
0216             if label[0] == "_": label = label[1:]   # seems labels starting "_" have special meaning to mpl, causing problems
0217             label = label.replace("solid","s")
0218
0219             if MODE == 2:
0220                 ax.scatter( gpos[:,H], gpos[:,V], s=SZ, color=color, label=label )
0221             elif MODE == 3:
0222                 pl.add_points( gpos[:,:3], color=color, label=label)
0223             pass
0224         pass
0225
0226         if MODE == 2:
0227             xlim, ylim = mpplt_focus_aspect()
0228             if not xlim is None:
0229                 ax.set_xlim(xlim)
0230                 ax.set_ylim(ylim)
0231             else:
0232                 log.info("mpplt_focus_aspect not enabled, use eg FOCUS=0,0,100 to enable ")
0233             pass
0234
0235             if 'POINT' in os.environ: # ADD POINT OR POINTS SPECIFIED BY TRIPLETS
0236                 point = efloatarray_("POINT", "0,0,0").reshape(-1,3)
0237                 ax.scatter( point[:,H], point[:,V], s=SZ, color="red", label="POINT" )
0238             pass
0239
0240             if 'LINE' in os.environ:
0241                 line = efloatarray_("LINE", "1,2,3,10,20,30,100,200,300").reshape(-1,3)
0242                 ax.plot( line[:,H], line[:,V], "ro-" )
0243             pass
0244
0245             if TIGHT:
0246                 ax.axis('off')
0247                 fig.tight_layout()
0248             pass
0249
0250         elif MODE == 3:
0251             pass
0252         pass
0253
0254
0255     def show(self, pl ):
0256         if MODE == 2:
0257             self.mp_show(pl)
0258         elif MODE == 3:
0259             self.pv_show(pl)
0260         pass
0261
0262     def pv_show(self, pl):
0263         cp = pl.show()
0264         print(cp)
0265
0266     def mp_show(self, pl):
0267
0268         fig, axs = pl
0269         assert len(axs) == 1
0270         ax = axs[0]
0271
0272         #ax.set_xlim(-60,60)  # use FOCUS=0,0,100
0273         #ax.set_ylim(-100,100)
0274
0275         locs = ["upper left","lower left", "upper right"]
0276         LOC = os.environ.get("LOC",locs[0])
0277         print("LOC : %s " % LOC)
0278         if LOC == "skip" or LOC == "":
0279             print("skip legend as LOC:[%s] " % LOC)
0280         else:
0281             ax.legend(loc=LOC,  markerscale=4)
0282         pass
0283
0284         thirdline = ""
0285         subtitle = ""
0286
0287         num_opp = len(self.opp)
0288
0289         if num_opp == 2:
0290             a,b = self.opp
0291             thirdline = a.label
0292             subtitle  = b.label
0293         else:
0294             ## eg "TO BT SA\n00 01 02"
0295             for one in self.opp:
0296                 print("one.label:%s " % one.label)
0297                 thirdline += one.label
0298             pass
0299         pass
0300
0301         subtitle = os.environ.get("SUBTITLE", "")
0302         print("mp_show num_opp:%d subtitle:%s " % (num_opp, subtitle) )
0303
0304         # adjust the position of the title, to legibly display 4 lines
0305         TOF = float(os.environ.get("TOF","0.99"))
0306
0307         #suptitle = "\n".join([TOPLINE,BOTLINE,thirdline])
0308         suptitle = TOPLINE
0309
0310
0311         fig.suptitle(suptitle, family="monospace", va="top", y=TOF, fontweight='bold' )
0312
0313         hv_thirdline = efloatarray_("HV_THIRDLINE","-0.05,1.02")
0314         hv_subtitle = efloatarray_("HV_SUBTITLE","1.05,-0.12")
0315         assert len(hv_thirdline) == 2
0316         assert len(hv_subtitle) == 2
0317
0318         ax.text( hv_thirdline[0],  hv_thirdline[1], thirdline, va='bottom', ha='left', family="monospace", fontsize=12, transform=ax.transAxes)
0319         ax.text( hv_subtitle[0],  hv_subtitle[1], subtitle, va='bottom', ha='right', family="monospace", fontsize=12, transform=ax.transAxes)
0320
0321         print("mp_show suptitle:[%s]" % suptitle )
0322         print("mp_show thirdline:[%s] HV_THIRDLINE:%s"  % (thirdline, str(hv_thirdline)) )
0323         print("mp_show subtitle:[%s]  HV_SUBTITLE:%s"   % (subtitle,  str(hv_subtitle)) )
0324
0325         fig.show()
0326
0327     def onephotonplot(self, pl, f):
0328         """
0329         :param pl:  plotting machinery
0330         :param f: expecting SEvt object with photon history,
0331                   NB pid must be set to +ve integer selecting the photon to plot anything
0332         """
0333         if f is None: return
0334         if f.pid < 0: return
0335
0336         self.opp.append(f)
0337
0338         if not hasattr(f,'r'): return
0339         if f.r is None: return
0340
0341         r = f.r
0342         off = f.off
0343
0344         rpos = r[:,0,:3] + off
0345
0346         if MODE == 2:
0347             fig, axs = pl
0348             assert len(axs) == 1
0349             ax = axs[0]
0350             if True:
0351                 mpplt_add_contiguous_line_segments(ax, rpos, axes=(H,V), label=None )
0352                 self.mp_plab(ax, f)
0353             pass
0354             if "nrm" in f.opt:
0355                 self.mp_a_normal(ax, f)
0356             pass
0357         elif MODE == 3:
0358             pass
0359             pvplt_add_contiguous_line_segments(pl, rpos )
0360         pass
0361
0362     def mp_plab(self, ax, f, backgroundcolor=None ):
0363         """
0364         point labels
0365         """
0366         r = f.r
0367         a = f.a
0368         ast = f.ast
0369         hv = r[:,0,(H,V)]
0370
0371         if backgroundcolor == None:
0372             backgroundcolor = os.environ.get("BGC", None)
0373         pass
0374         for i in range(len(r)):
0375
0376             if "idx" in f.opt:
0377                 plab = str(i)
0378             elif "ast" in f.opt:
0379                 plab = chr(ast[i])
0380             else:
0381                 plab = None
0382             pass
0383             if plab is None: continue
0384
0385             dx,dy = 0,0
0386
0387             if "pdy" in f.opt: dy = 10
0388             if "ndy" in f.opt: dy = -10
0389
0390             print("mp_plab: f.opt %s %10.4f %10.4f " % (f.opt, dx, dy))
0391
0392             if backgroundcolor is None:
0393                 ax.text(dx+hv[i,0],dy+hv[i,1], plab, fontsize=15 )
0394             else:
0395                 ax.text(dx+hv[i,0],dy+hv[i,1], plab, fontsize=15, backgroundcolor=backgroundcolor )
0396             pass
0397             ## POINT-BY-POINT LABELS
0398         pass
0399
0400
0401
0402     def mp_a_normal(self, ax, f):
0403         """
0404         Access customBoundaryStatus char
0405
0406         t.aux[261,:32,1,3].copy().view(np.int8)[::4].copy().view("|S32")
0407
0408         """
0409         if not hasattr(f,'r'): return
0410         if not hasattr(f,'a'): return
0411         if f.r is None: return
0412
0413         r = f.r
0414         a = f.a
0415         sc = float(os.environ.get("NRM","80"))
0416         #cbs = a[:len(a),1,3].copy().view(np.int8)[::4]
0417         #assert len(cbs) == len(a)
0418
0419         ast = f.ast
0420         assert len(ast) == len(a) == len(r)
0421
0422
0423         for i in range(len(a)):
0424             nrm = a[i,3,:3]
0425             pos = r[i,0,:3]
0426             npos = pos - sc*nrm/2
0427             #mpplt_add_line(ax, pos-sc*nrm, pos+sc*nrm, AXES )
0428             ax.arrow( npos[H], npos[V], sc*nrm[H], sc*nrm[V], head_width=10, head_length=10, fc='k', ec='k' )
0429         pass
0430
0431
0432
0433 if __name__ == '__main__':
0434     logging.basicConfig(level=logging.INFO)
0435
0436     fold = SFOLD if not SFOLD is None else FOLD
0437
0438     log.info(" -- U4SimtraceSimpleTest.Load FOLD/SFOLD fold:%s" % fold  )
0439     s = U4SimtraceSimpleTest.Load(fold, symbol="s")    # mandatory first geometry
0440     assert not s is None
0441
0442
0443     log.info(" -- U4SimtraceSimpleTest.Load TFOLD" )
0444     t = U4SimtraceSimpleTest.Load(TFOLD, symbol="t")   # optional second geometry
0445     ## SFOLD, TFOLD and s, t correspond to intersection geometries from U4SimtraceTest
0446
0447
0448     log.info(" -- SEvt.Load AFOLD" )
0449     a = SEvt.Load(AFOLD, symbol="a")   # optional photon histories
0450     log.info(" -- SEvt.Load BFOLD" )
0451     b = SEvt.Load(BFOLD, symbol="b")
0452     ## AFOLD, BFOLD and a, b are photon histories from SEvt
0453
0454     if not a is None: print('a',a.qtab_,"\n",a.qtab)
0455     if not b is None: print('b',b.qtab_,"\n",b.qtab)
0456
0457
0458     pl = plotter(label="U4SimtraceSimpleTest.py")  # MODE:2 (fig,ax)  MODE:3 pv plotter
0459
0460     if not pl is None:
0461         s.geoplot(pl)
0462         if not t is None: t.geoplot(pl)
0463
0464         s.onephotonplot(pl, a)   # must set APID to int index of photon for this to do anything
0465         s.onephotonplot(pl, b)   # must set BPID to int index of photon for this to do anything
0466
0467         s.show(pl)
0468     pass
0469
0470
0471
0472 pass
0473
0474
0475