eic-opticks/sysrap/sframe.py

0001 #!/usr/bin/env python
0002
0003 import logging
0004 log = logging.getLogger(__name__)
0005
0006 from opticks.ana.npmeta import NPMeta
0007 from opticks.ana.axes import *
0008 from opticks.ana.pvplt import mpplt_focus, SIZE
0009
0010 import os, numpy as np
0011 eary_ = lambda ekey, edef:np.array( list(map(float, os.environ.get(ekey,edef).split(","))) )
0012
0013 from opticks.ana.eget import efloat_
0014
0015
0016 class sframe(object):
0017     @classmethod
0018     def Load(cls, fold=None, name="sframe.npy"):
0019         if fold is None:
0020             fold = os.environ.get("FOLD", "")
0021         pass
0022         if fold.endswith(name):
0023             path = fold
0024         else:
0025             path = os.path.join(fold, name)
0026         pass
0027         return cls(path)
0028
0029
0030     @classmethod
0031     def DetermineAxes(cls, nx, ny, nz):
0032         """
0033         :param nx:
0034         :param nx:
0035         :param nx:
0036
0037         With planar axes the order is arranged to make the longer axis the first horizontal one
0038         followed by the shorter axis as the vertical one.
0039
0040             +------+
0041             |      |  nz     ->  ( Y, Z )    ny_over_nz > 1
0042             +------+
0043                ny
0044
0045             +------+
0046             |      |  ny     ->  ( Z, Y )    ny_over_nz < 1
0047             +------+
0048                nz
0049
0050         """
0051         if nx == 0 and ny > 0 and nz > 0:
0052             ny_over_nz = float(ny)/float(nz)
0053             axes = (Y,Z) if ny_over_nz > 1 else (Z,Y)
0054         elif nx > 0 and ny == 0 and nz > 0:
0055             nx_over_nz = float(nx)/float(nz)
0056             axes = (X,Z) if nx_over_nz > 1 else (Z,X)
0057         elif nx > 0 and ny > 0 and nz == 0:
0058             nx_over_ny = float(nx)/float(ny)
0059             axes = (X,Y) if nx_over_ny > 1 else (Y,X)
0060         else:
0061             axes = (X,Y,Z)
0062         pass
0063         return axes
0064
0065
0066     @classmethod
0067     def CombineFrame(cls, framelist_):
0068         """
0069         HMM: probably should do more consistency checks between all the frames
0070         """
0071         framelist = list(filter(None, framelist_))
0072         axes = None
0073         for f in framelist:
0074             if axes is None:
0075                 axes = f.axes
0076             else:
0077                 assert axes == f.axes
0078             pass
0079         pass
0080         return framelist[0] if len(framelist) > 0 else None
0081
0082
0083     def __init__(self, path, clear_identity=True ):
0084         """
0085         Whether clear_identity makes any material difference depends on the identity values.
0086         But it should be done anyhow.  For some identity values they will appear as nan in float.
0087         """
0088         if path is None:
0089             return
0090         pass
0091
0092         metapath = path.replace(".npy", "_meta.txt")
0093         if os.path.exists(metapath):
0094             meta = NPMeta.Load(metapath)
0095         else:
0096             meta = None
0097         pass
0098         self.meta = meta
0099
0100         a = np.load(path)
0101         i = a.view(np.int32)
0102
0103         self.shape = a.shape  # for fold to treat like np.array
0104         self.path = path
0105         self.a = a
0106         self.i = i
0107
0108         ce = a[0,0]
0109         ix0,ix1,iy0,iy1 = i[0,1,:4]        # q1.i.xyzw
0110         iz0,iz1,num_photon = i[0,2,:3]     # q2.i.xyz
0111         gridscale = a[0,2,3]               # q2.f.w
0112
0113
0114         midx, mord, gord = i[0,3,:3]       # q3.i.xyz
0115         inst = i[0,3,3]     # q3.i.w
0116
0117
0118
0119         self.midx = midx
0120         self.mord = mord
0121         self.gord = gord
0122         self.inst = inst
0123
0124         ## replacing ana/gridspec.py:GridSpec
0125
0126
0127
0128         grid = "".join(["ix0 %(ix0)4d ix1 %(ix1)4d ",
0129                         "iy0 %(iy0)4d iy1 %(iy1)4d ",
0130                         "iz0 %(iz0)4d iz1 %(iz1)4d ",
0131                         "num_photon %(num_photon)4d gridscale %(gridscale)10.4f"]) % locals()
0132
0133
0134
0135
0136
0137
0138         self.grid = grid
0139         self.ce = ce
0140         self.ix0 = ix0
0141         self.ix1 = ix1
0142         self.iy0 = iy0
0143         self.iy1 = iy1
0144         self.iz0 = iz0
0145         self.iz1 = iz1
0146         self.num_photon = num_photon
0147         self.gridscale = gridscale
0148
0149
0150
0151         target = "midx %(midx)6d mord %(mord)6d gord %(gord)6d       inst %(inst)7d   " % locals()
0152         self.target = target
0153
0154
0155         ins_idx_1 = i[1,0,3] - 1   # q0.u.w-1  see sqat4.h qat4::getIdentity qat4::setIdentity
0156         gas_idx_1 = i[1,1,3] - 1   # q1.u.w-1
0157         ias_idx_1 = i[1,2,3] - 1   # q2.u.w-1
0158
0159         ins_idx_2 = i[2,0,3] - 1   # q0.u.w-1  see sqat4.h qat4::getIdentity qat4::setIdentity
0160         gas_idx_2 = i[2,1,3] - 1   # q1.u.w-1
0161         ias_idx_2 = i[2,2,3] - 1   # q2.u.w-1
0162
0163         ins_idx_match = ins_idx_1 == ins_idx_2
0164         gas_idx_match = gas_idx_1 == gas_idx_2
0165         ias_idx_match = ias_idx_1 == ias_idx_2
0166
0167         if not ins_idx_match: log.warning(f" ins_idx_match {ins_idx_match} ins_idx_1 {ins_idx_1} ins_idx_2 {ins_idx_2} ")
0168         if not gas_idx_match: log.warning(f" gas_idx_match {gas_idx_match} gas_idx_1 {gas_idx_1} gas_idx_2 {gas_idx_2} ")
0169         if not ias_idx_match: log.warning(f" ias_idx_match {ias_idx_match} ias_idx_1 {ias_idx_1} ias_idx_2 {ias_idx_2} ")
0170
0171         #assert ins_idx_match
0172         #assert gas_idx_match
0173         #assert ias_idx_match
0174
0175         ins_idx = ins_idx_1
0176         gas_idx = gas_idx_1
0177         ias_idx = ias_idx_1
0178         qat4id = "ins_idx %(ins_idx)6d gas_idx %(gas_idx)4d %(ias_idx)4d " % locals()
0179
0180         self.ins_idx = ins_idx
0181         self.gas_idx = gas_idx
0182         self.ias_idx = ias_idx
0183         self.qat4id = qat4id
0184
0185         m2w = a[1].copy()
0186         w2m = a[2].copy()
0187
0188         if clear_identity:
0189             m2w[0,3] = 0.
0190             m2w[1,3] = 0.
0191             m2w[2,3] = 0.
0192             m2w[3,3] = 1.
0193
0194             w2m[0,3] = 0.
0195             w2m[1,3] = 0.
0196             w2m[2,3] = 0.
0197             w2m[3,3] = 1.
0198         pass
0199
0200         self.m2w = m2w
0201         self.w2m = w2m
0202         self.id = np.dot( m2w, w2m )
0203
0204
0205         ins = i[3,0,0]   # aux.q0.i.x
0206         gas = i[3,0,1]
0207         ias = i[3,0,2]
0208
0209         propagate_epsilon = a[3,1,0]   # aux.q1.f.x
0210
0211
0212         ins_gas_ias = " ins %(ins)6d gas %(gas)4d ias %(ias)4d " % locals()
0213
0214         self.ins = ins
0215         self.gas = gas
0216         self.ias = ias
0217         self.ins_gas_ias = ins_gas_ias
0218         self.propagate_epsilon = propagate_epsilon
0219
0220         self.init_grid()
0221         self.init_view()
0222
0223
0224     def init_grid(self):
0225         """
0226         replacing ana/gridspec.py
0227         """
0228         gord = self.gord
0229         ix0 = self.ix0
0230         ix1 = self.ix1
0231         iy0 = self.iy0
0232         iy1 = self.iy1
0233         iz0 = self.iz0
0234         iz1 = self.iz1
0235         gridscale = self.gridscale
0236         ce = self.ce
0237         extent = ce[3]
0238         s_extent = "e:%7.3f" % extent
0239
0240         nx = (ix1 - ix0)//2
0241         ny = (iy1 - iy0)//2
0242         nz = (iz1 - iz0)//2
0243         coords = "RTP" if gord == -3 else "XYZ"   ## NB RTP IS CORRECT ORDERING radiusUnitVec:thetaUnitVec:phiUnitVec
0244         axes = self.DetermineAxes(nx, ny, nz)
0245         other_axis = Axes.OtherAxis(axes)
0246         planar = len(axes) == 2
0247
0248         bbox = np.zeros( (2,3), dtype=np.float32 )
0249         bbox[0] = list(map(float,[ix0,iy0,iz0]))
0250         bbox[1] = list(map(float,[ix1,iy1,iz1]))
0251         bbox *= gridscale*ce[3]
0252
0253
0254         if planar:
0255             H, V = axes
0256             axlabels =  coords[H], coords[V]
0257             s_bbox = "bb "
0258             s_bbox += "  %7.4f %7.4f " % tuple(bbox[0,axes])
0259             s_bbox += "  %7.4f %7.4f " % tuple(bbox[1,axes])
0260         else:
0261             H, V, D = axes
0262             axlabels =  coords[H], coords[V], coords[D]
0263             s_bbox = "bb %7.4f %7.4f %7.4f    %7.4f %7.4f %7.4f " % tuple(bbox.ravel())
0264         pass
0265
0266
0267         self.nx = nx
0268         self.ny = ny
0269         self.nz = nz
0270         self.coords = coords
0271         self.axes = axes
0272         self.other_axis = other_axis
0273         self.axlabels = axlabels
0274
0275         self.bbox = bbox  # without ce_offset
0276         self.s_bbox = s_bbox
0277
0278         self.extent = extent
0279         self.s_extent = s_extent
0280
0281
0282
0283     def init_view(self):
0284         """
0285         HMM: input EYE is ignored for planar
0286         """
0287         ce = self.ce
0288         axes = self.axes
0289         coords = self.coords
0290         planar = len(axes) == 2
0291
0292         # below default from envvars are overridden for planar data
0293         eye = eary_("EYE","1.,1.,1.")
0294         up  = eary_("UP","0.,0.,1.")
0295         off  = eary_("OFF","0.,0.,1.")
0296         look = eary_("LOOK","0.,0.,0.")
0297         #look = ce[:3]
0298
0299         EYES = efloat_("EYES", "6.")   # TODO: why 6 ? how to control FOV to gain more control of this
0300         extent = ce[3]
0301
0302         if planar:
0303             H, V = axes
0304             up  = Axes.Up(H,V)
0305             off = Axes.Off(H,V)
0306             eye = look + extent*off*EYES
0307         else:
0308             H, V, D = axes
0309             axlabels =  coords[H], coords[V], coords[D]
0310
0311             up = XYZ.up
0312             off = XYZ.off
0313             ## hmm in 3D case makes less sense : better to just use the input EYE
0314
0315             eye = ce[3]*eye*EYES
0316         pass
0317
0318         self.look = look
0319         self.up = up
0320         self.off = off
0321         self.eye = eye
0322         self.thirdline = self.s_bbox + self.s_extent
0323
0324
0325     def pv_compose(self, pl, local=True):
0326         """
0327         #pl.view_xz()   ## TODO: see if view_xz is doing anything when subsequently set_focus/viewup/position
0328         """
0329
0330         RESET = "RESET" in os.environ
0331         PARA = "PARA" in os.environ
0332         ZOOM = efloat_("ZOOM", "1.")
0333         PVGRID = "PVGRID" in os.environ
0334
0335         if PARA:
0336             pl.camera.ParallelProjectionOn()
0337         pass
0338
0339         look = self.look if local else self.ce[:3]   ## HMM: same ?
0340         eye = look + self.off
0341         up = self.up
0342
0343         ## for reset=True to succeed to auto-set the view, must do this after add_points etc..
0344
0345         #eye = look + self.off
0346
0347         look = self.look
0348         eye = self.eye
0349         up = self.up
0350
0351         log.info("frame.pv_compose look:%s eye: %s up:%s  PARA:%s RESET:%d ZOOM:%s  " % (str(look), str(eye), str(up), RESET, PARA, ZOOM ))
0352
0353         pl.set_focus(    look )
0354         pl.set_viewup(   up )
0355         pl.set_position( eye, reset=RESET )   ## for reset=True to succeed to auto-set the view, must do this after add_points etc..
0356         pl.camera.Zoom(ZOOM)
0357
0358         if PVGRID:
0359             pl.show_grid()
0360         pass
0361
0362
0363
0364     def __repr__(self):
0365
0366         l_ = lambda k,v:"%-12s : %s" % (k, v)
0367
0368         return "\n".join(
0369                   [
0370                     l_("sframe",""),
0371                     l_("path",self.path),
0372                     l_("meta",repr(self.meta)),
0373                     l_("ce", repr(self.ce)),
0374                     l_("grid", self.grid),
0375                     l_("bbox", repr(self.bbox)),
0376                     l_("target", self.target),
0377                     l_("qat4id", self.qat4id),
0378                     l_("m2w",""), repr(self.m2w), "",
0379                     l_("w2m",""), repr(self.w2m), "",
0380                     l_("id",""),  repr(self.id) ,
0381                     l_("ins_gas_ias",self.ins_gas_ias)
0382                    ])
0383
0384
0385     @classmethod
0386     def FakeXZ(cls, e=10):
0387         path = None
0388         fr = cls(path)
0389         fr.axes = X,Z
0390         fr.axlabels = "X","Z"
0391         fr.bbox = np.array( [[-e,-e,-e],[e,e,e]] )
0392         return fr
0393
0394
0395     def mp_subplots(self, mp):
0396         pass
0397         H,V = self.axes
0398         _H,_V = self.axlabels
0399
0400         xlim = self.bbox[:,H]
0401         ylim = self.bbox[:,V]
0402
0403         log.info("mp_subplots xlim %s ylim %s " % (str(xlim), str(ylim)))
0404         #xlim, ylim = mpplt_focus(xlim, ylim)
0405
0406         topline = os.environ.get("TOPLINE", "sframe.py:mp_subplots:TOPLINE")
0407         botline = os.environ.get("BOTLINE", "sframe.py:mp_subplots:BOTLINE")
0408         thirdline = getattr(self, "thirdline", "thirdline" )
0409
0410         title = [topline, botline, thirdline]
0411
0412         fig, ax = mp.subplots(figsize=SIZE/100.)
0413         fig.suptitle("\n".join(title))
0414
0415         ax.set_aspect('equal')
0416         ax.set_xlim(xlim)
0417         ax.set_ylim(ylim)
0418         ax.set_xlabel(_H)
0419         ax.set_ylabel(_V)
0420
0421         self.fig = fig
0422         self.ax = ax
0423
0424         return fig, ax
0425
0426     def mp_scatter(self, pos, **kwa):
0427         assert pos.ndim == 2 and pos.shape[1] == 3
0428         H,V = self.axes
0429         self.ax.scatter( pos[:,H], pos[:,V], **kwa)
0430
0431     def mp_arrow(self, pos, mom, **kwa):
0432         assert pos.ndim == 2 and pos.shape[1] == 3
0433         assert mom.ndim == 2 and mom.shape[1] == 3
0434         assert len(pos) == len(mom)
0435         H,V = self.axes
0436
0437         for i in range(len(pos)):
0438             self.ax.arrow( pos[i,H], pos[i,V], mom[i,H], mom[i,V] )
0439         pass
0440
0441
0442     def mp_legend(self):
0443         NOLEGEND = "NOLEGEND" in os.environ
0444         LSIZ = float(os.environ.get("LSIZ",50))
0445         LOC = os.environ.get("LOC", "upper right")
0446         log.info("mp_legend LSIZ %s LOC %s NOLEGEND %s " % (LSIZ,LOC,NOLEGEND) )
0447         if not NOLEGEND:
0448             lgnd = self.ax.legend(loc=LOC)
0449             for handle in lgnd.legendHandles:
0450                 handle.set_sizes([LSIZ])
0451             pass
0452         pass
0453
0454
0455
0456
0457 if __name__ == '__main__':
0458
0459     fr = sframe.Load()
0460     print(fr)
0461
0462
0463