import numpy as np
def RotX(th):
sth, cth = np.sin(th), np.cos(th)
return np.array([[1,0,0],[0,cth,-sth],[0,sth,cth]])
def RotY(th):
sth, cth = np.sin(th), np.cos(th)
return np.array([[cth,0,sth],[0,1,0],[-sth,0,cth]])
def RotZ(th):
sth, cth = np.sin(th), np.cos(th)
return np.array([[cth,-sth,0],[sth,cth,0],[0,0,1]])
[docs]class CoordSys:
"""A coordinate system against which to measure surfaces or rays.
Coordinate systems consist of an origin and a rotation. The ``origin``
attribute specifies where in 3D space the current coordinate system's
origin lands in the global coordinate system. The rotation ``rot``
specifies the 3D rotation matrix to apply to the global coordinate axes to
yield the axes of the this coordinate system.
Parameters
----------
origin : ndarray of float, shape (3,)
Origin of coordinate system in global coordinates.
rot : ndarray of float, shape (3, 3)
Rotation matrix taking global axes into current system axes.
"""
def __init__(self, origin=None, rot=None):
if origin is None:
origin = np.zeros(3, dtype=float)
if rot is None:
rot = np.eye(3, dtype=float)
self.origin = np.array(origin)
self.rot = np.array(rot)
@property
def xhat(self):
"""ndarray of float, shape (3,): Orientation of local x vector in
global coordinates.
"""
return self.rot[:, 0]
@property
def yhat(self):
"""ndarray of float, shape (3,): Orientation of local y vector in
global coordinates.
"""
return self.rot[:, 1]
@property
def zhat(self):
"""ndarray of float, shape (3,): Orientation of local z vector in
global coordinates.
"""
return self.rot[:, 2]
[docs] def shiftGlobal(self, dr):
"""Return new CoordSys with origin shifted along global axes.
Parameters
----------
dr : ndarray of float, shape (3,)
Amount to shift in meters.
Returns
-------
CoordSys
"""
return CoordSys(self.origin+dr, self.rot)
[docs] def shiftLocal(self, dr):
"""Return new CoordSys with origin shifted along local axes.
Parameters
----------
dr : ndarray of float, shape (3,)
Amount to shift in meters.
Returns
-------
CoordSys
"""
# Rotate the shift into global coordinates, then do the shift globally
return self.shiftGlobal(self.rot@dr)
[docs] def toGlobal(self, v):
"""Convert vector(s) from local to global coordinates.
Parameters
----------
v : ndarray of float, shape (n, 3)
Vector(s) in local coordinates.
Returns
-------
vv : ndarray of float, shape (n, 3)
Vector in global coordinates.
"""
v = np.array(v, dtype=float)
return (self.rot@v.T).T + self.origin
[docs] def toLocal(self, v):
"""Convert vector(s) from global to local coordinates.
Parameters
----------
v : ndarray of float, shape (n, 3)
Vector(s) in global coordinates.
Returns
-------
vv : ndarray of float, shape (n, 3)
Vector in local coordinates.
"""
v = np.array(v, dtype=float)
v -= self.origin
return (self.rot.T@v.T).T
[docs] def rotateGlobal(self, rot, rotCenter=(0,0,0), coordSys=None):
"""Return new CoordSys rotated with respect to global axes.
Parameters
----------
rot : ndarray of float, shape (3, 3)
Rotation matrix to apply.
rotCenter : ndarray of float, shape (3,)
Point about which to rotate.
coordSys : CoordSys
Coordinate system in which rotCenter is specified.
Returns
-------
CoordSys
"""
if coordSys is None:
coordSys = CoordSys()
# Find rot center in global coordinates
globalRotCenter = coordSys.toGlobal(rotCenter)
# Then rotate about this center
return CoordSys(
rot@(self.origin-globalRotCenter)+globalRotCenter,
rot@self.rot
)
[docs] def rotateLocal(self, rot, rotCenter=(0,0,0), coordSys=None):
"""Return new CoordSys rotated with respect to local axes.
Parameters
----------
rot : ndarray of float, shape (3, 3)
Rotation matrix to apply.
rotCenter : ndarray of float, shape (3,)
Point about which to rotate.
coordSys : CoordSys
Coordinate system in which rotCenter is specified.
Returns
-------
CoordSys
"""
if coordSys is None:
coordSys = self
# Find rot center in global coordinates
globalRotCenter = coordSys.toGlobal(rotCenter)
# first rotate rot into global coords: (self.rot rot self.rot.T),
# then apply that: (self.rot rot self.rot.T) self.rot = self.rot rot
rTmp = self.rot@rot
return CoordSys(
rTmp@self.rot.T@(self.origin-globalRotCenter)+globalRotCenter,
rTmp
)
def __getstate__(self):
return self.origin, self.rot
def __setstate__(self, d):
self.origin, self.rot = d
def __eq__(self, rhs):
if not isinstance(rhs, CoordSys): return False
return (
np.array_equal(self.origin, rhs.origin) and
np.array_equal(self.rot, rhs.rot)
)
def __ne__(self, rhs):
return not (self == rhs)
def __hash__(self):
return hash((
"batoid.CoordSys",
tuple(self.origin.tolist()),
tuple(self.rot.ravel().tolist())
))
def copy(self):
return CoordSys(self.origin, self.rot)
def __repr__(self):
rotstr = np.array2string(self.rot, separator=', ').replace('\n', '')
return f"CoordSys({self.origin!r}, array({rotstr}))"