eic-opticks/sysrap/smath.py

0001 #!/usr/bin/env python
0002
0003 import numpy as np
0004
0005 def rotateUz(d, u):
0006     """
0007     :param d:  array of vectors with shape (n,3)
0008     :param u:  array of single vector with shape (3,)
0009     :return d1: array of vectors with shape (n,3)
0010
0011     NumPy translation of smath.h rotateUz::
0012
0013         inline SMATH_METHOD void smath::rotateUz(float3& d, const float3& u )
0014         {
0015             float up = u.x*u.x + u.y*u.y ;
0016             if (up>0.f)
0017             {
0018                 up = sqrt(up);
0019                 float px = d.x ;
0020                 float py = d.y ;
0021                 float pz = d.z ;
0022                 d.x = (u.x*u.z*px - u.y*py)/up + u.x*pz;
0023                 d.y = (u.y*u.z*px + u.x*py)/up + u.y*pz;
0024                 d.z =    -up*px +                u.z*pz;
0025             }
0026             else if (u.z < 0.f )
0027             {
0028                 d.x = -d.x;
0029                 d.z = -d.z;
0030             }
0031         }
0032
0033     """
0034     assert u.shape == (3,)
0035     assert d.shape == (len(d),3)
0036     d1 = np.zeros( (len(d),3 ), dtype=d.dtype )
0037
0038     up =  u[0]*u[0] + u[1]*u[1]
0039     if up > 0.:
0040        up = np.sqrt(up)
0041        px = d[:,0]
0042        py = d[:,1]
0043        pz = d[:,2]
0044        d1[:,0] = (u[0]*u[2]*px - u[1]*py)/up + u[0]*pz
0045        d1[:,1] = (u[1]*u[2]*px + u[0]*py)/up + u[1]*pz
0046        d1[:,2] = -up*px + u[2]*pz
0047     elif u[2] < 0.:
0048        d1[:,0] = -d[:,0]
0049        d1[:,1] =  d[:,1]  # diff as returning not inplace changing
0050        d1[:,2] = -d[:,2]
0051     pass
0052     return d1
0053
0054
0055 def rotateUz_(d, u):
0056     """
0057     :param d:  array of single vector with shape (3,)
0058     :param u:  array of vectors with shape (n,3)
0059     :return d1: array of vectors with shape (n,3)
0060
0061     THIS NEEDS MORE TESTING
0062     """
0063     assert u.shape == (len(u),3)
0064     assert d.shape == (3,)
0065     d1 = np.zeros( (len(u),3 ), dtype=d.dtype )
0066
0067     _up =  u[:,0]*u[:,0] + u[:,1]*u[:,1]
0068     wp = np.where( _up > 0. )
0069     wn = np.where( np.logical_and( _up <= 0., u[:,2] < 0. ))
0070
0071     up = np.zeros( (len(u),), dtype=d.dtype )
0072     up[wp] = np.sqrt(_up)
0073     px = d[0]
0074     py = d[1]
0075     pz = d[2]
0076
0077     d1[wp,0] = (u[wp,0]*u[wp,2]*px - u[wp,1]*py)/up + u[wp,0]*pz
0078     d1[wp,1] = (u[wp,1]*u[wp,2]*px + u[wp,0]*py)/up + u[wp,1]*pz
0079     d1[wp,2] = -up[wp]*px + u[wp,2]*pz
0080
0081     d1[wn,0] = -d[0]
0082     d1[wn,1] =  d[1]
0083     d1[wn,2] = -d[2]
0084
0085     return d1
0086
0087
0088 def test_rotateUz():
0089     s = np.sqrt(0.5, dtype=np.float64)
0090     u = np.array([s,0,-s], dtype=s.dtype)
0091     d = np.array([[1,0,0],[s,s,0],[0,1,0],[-s,s,0],[-1,0,0],[-s,-s,0],[0,-1,0],[s,-s,0],[1,0,0]], dtype=s.dtype )
0092     d1 = rotateUz( d, u )
0093
0094     print("u\n",u)
0095     print("d\n",d)
0096     print("d1\n",d1)
0097
0098
0099
0100 if __name__ == '__main__':
0101
0102     s = np.sqrt(0.5, dtype=np.float64)
0103     u = np.array([[1,0,0],[s,s,0],[0,1,0],[-s,s,0],[-1,0,0],[-s,-s,0],[0,-1,0],[s,-s,0],[1,0,0]], dtype=s.dtype )
0104     d = np.array([s,0,-s], dtype=s.dtype)
0105     d1 = rotateUz_( d, u )
0106
0107     print("u\n",u)
0108     print("d\n",d)
0109     print("d1\n",d1)
0110
0111