eic-opticks/ana/parallel_input_photons.py

0001 #!/usr/bin/env python
0002 """
0003
0004
0005 Repeat input photons and offset them
0006 using the polarization direction.::
0007
0008
0009           r = 5
0010           j0 = -r//2
0011           j1 =  r//2
0012
0013
0014                        ^  direction
0015            |  |  |  |  |
0016            |  |  |  |  |
0017            |  |  |  |  |
0018            |  |  |  |  |
0019            |  |  |  |  |
0020        -------------------------------> polarization
0021
0022
0023     photon item (4,4)
0024
0025       0: pos_x, pos_y, pos_z, time
0026       1: dir_x, dir_y, dir_z, weight
0027       2: pol_x, pol_y, pol_z, wavelength
0028       4: flags
0029
0030
0031
0032
0033 """
0034 import numpy as np
0035 from opticks.ana.input_photons import InputPhotons
0036 import matplotlib.pyplot as plt
0037 import pyvista as pv
0038
0039
0040 def mock_photons(o):
0041     p = np.zeros([o,4,4], dtype=np.float32)
0042     for k in range(o):
0043         p[k,2,:3] = [0,0,1]
0044     pass
0045     return p
0046
0047
0048 def mplt(pp):
0049
0050     pos_x, pos_y, pos_z = pp[:,0,0], pp[:,0,1], pp[:,0,2]
0051     dir_x, dir_y, dir_z = pp[:,1,0], pp[:,1,1], pp[:,1,2]
0052
0053     plt.ion()
0054     fig = plt.figure(figsize=(6,5.5))
0055     ax = fig.add_subplot(projection='3d')
0056     ax.scatter( pos_x, pos_y, pos_z )
0057     ax.scatter( pos_x+dir_x, pos_y+dir_y, pos_z+dir_z )
0058     fig.show()
0059
0060 def pvplt(pp):
0061     pl = pv.Plotter()
0062
0063     pos = pp[:,0,:3]
0064     dir = pp[:,1,:3]
0065     pol = pp[:,2,:3]
0066     oth = np.cross( pol, dir )
0067     mag = 1
0068     pl.add_arrows( pos, dir, mag=mag*1.0,  color='#FF0000', point_size=2.0 )
0069     pl.add_arrows( pos, pol, mag=mag*1.0,  color='#00FF00', point_size=2.0 )
0070     pl.add_arrows( pos, oth, mag=mag*1.0,  color='#0000FF', point_size=2.0 )
0071
0072     pl.show_grid()
0073     cpos = pl.show()
0074
0075
0076 def test_parallelise_1d(p):
0077     r = 10
0078     pp = InputPhotons.Parallelize1D(p, r, offset=True)
0079     pvplt(pp)
0080     return pp
0081
0082 def test_parallelise_2d(p):
0083     rr = [10,10]
0084     pp = InputPhotons.Parallelize2D(p, rr, offset=True)
0085     pvplt(pp)
0086     return pp
0087
0088
0089 if __name__ == '__main__':
0090     p = InputPhotons.GenerateCubeCorners()
0091     #p = InputPhotons.GenerateAxes()
0092     #p = p[0:1]
0093
0094     pp = test_parallelise_2d(p)
0095