sysrap/tests/SGenerate__test.py

0001 #!/usr/bin/env python
0002 """
0003 SGenerate__test.py
0004 ====================
0005
0006 1. load f from $FOLD
0007 2. check f.gs f.ph shapes
0008 3. for MODE=3 make 3D plot of generated photons
0009
0010 ::
0011
0012     hookup_conda_ok  # BYOC
0013     MODE=3 ~/o/sysrap/tests/SGenerate__test.sh pdb
0014
0015
0016 """
0017 import os, numpy as np
0018 from opticks.ana.fold import Fold
0019
0020 MODE = int(os.environ.get("MODE",0))
0021
0022 if MODE in [2,3]:
0023     import opticks.ana.pvplt as pvp
0024 else:
0025     pvp = None
0026 pass
0027
0028 if __name__ == '__main__':
0029     f = Fold.Load(symbol="f")
0030     print(repr(f))
0031
0032
0033     GS_NAME = os.environ.get("SGenerate__test_GS_NAME", "gs.npy").replace(".npy", "")
0034     PH_NAME = os.environ.get("SGenerate__test_PH_NAME", "ph.npy").replace(".npy", "")
0035
0036     gs = getattr(f, GS_NAME)
0037     ph = getattr(f, PH_NAME)
0038
0039     print("GS_NAME %s gs.shape %s\n" % (GS_NAME, str(gs.shape)))
0040     print("PH_NAME %s ph.shape %s\n" % (PH_NAME, str(ph.shape)))
0041
0042     pos = ph[:,0,:3]
0043     mom = ph[:,1,:3]
0044     pol = ph[:,2,:3]
0045
0046     if MODE == 3:
0047         pl = pvp.pvplt_plotter()
0048         #pl.add_points(pos, point_size=20)
0049         #pvplt_arrows(pl, pos, mom, factor=20 )
0050
0051         pvp.pvplt_polarized(pl, pos, mom, pol, factor=20  )
0052
0053         pl.show()
0054     pass