Source code for batoid.coordSys

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}))"