Make matrix, vector multiply functions public

This commit is contained in:
Brandon Rhodes 2020-06-12 14:18:08 -04:00
parent 56f1ada71f
commit 7e2fc1919f
6 changed files with 35 additions and 33 deletions

View File

@ -20,18 +20,20 @@ def _T(M):
"""Swap the first two dimensions of an array."""
return rollaxis(M, 1)
def _mxv(M, v):
"""Multiply an NxN matrix by a vector."""
def mxv(M, v):
"""Matrix times vector: multiply an NxN matrix by a vector."""
return einsum('ij...,j...->i...', M, v)
def _mxm(M1, M2):
"""Multiply two NxN matrices together."""
def mxm(M1, M2):
"""Matrix times matrix: multiply two NxN matrices."""
return einsum('ij...,jk...->ik...', M1, M2)
def _mxmxm(M1, M2, M3):
"""Multiply three NxN matrices together."""
def mxmxm(M1, M2, M3):
"""Matrix times matrix times matrix: multiply 3 NxN matrices together."""
return einsum('ij...,jk...,kl...->il...', M1, M2, M3)
_mxv, _mxm, _mxmxm = mxv, mxm, mxmxm # In case anyone imported old name
def length_of(xyz):
"""Given a 3-element array `[x y z]`, return its length.

View File

@ -4,7 +4,7 @@ import re
from numpy import array, cos, nan, sin
from jplephem.pck import DAF, PCK
from .constants import ASEC2RAD, AU_KM, DAY_S, tau
from .functions import _T, _mxv, _mxm, _mxmxm, rot_x, rot_y, rot_z
from .functions import _T, mxv, mxm, mxmxm, rot_x, rot_y, rot_z
from .units import Angle, Distance
from .vectorlib import VectorFunction
@ -105,7 +105,7 @@ class PlanetaryConstants(object):
matrix.shape = 3, 3
for angle, axis in list(zip(angles, axes)):
rot = _rotations[axis]
matrix = _mxm(rot(angle * scale), matrix)
matrix = mxm(rot(angle * scale), matrix)
elif spec == 'MATRIX':
matrix = self.assignments['TKFRAME_{0}_MATRIX'.format(integer)]
matrix = array(matrix)
@ -159,16 +159,16 @@ class Frame(object):
def rotation_at(self, t):
"""Return the rotation matrix for this frame at time ``t``."""
ra, dec, w = self._segment.compute(t.tdb, 0.0, False)
R = _mxm(rot_z(-w), _mxm(rot_x(-dec), rot_z(-ra)))
R = mxm(rot_z(-w), mxm(rot_x(-dec), rot_z(-ra)))
if self._matrix is not None:
R = _mxm(self._matrix, R)
R = mxm(self._matrix, R)
return R
def rotation_and_rate_at(self, t):
"""Return rotation and rate matrices for this frame at time ``t``."""
components, rates = self._segment.compute(t.tdb, 0.0, True)
ra, dec, w = components
R = _mxm(rot_z(-w), _mxm(rot_x(-dec), rot_z(-ra)))
R = mxm(rot_z(-w), mxm(rot_x(-dec), rot_z(-ra)))
zero = w * 0.0
one = 1.0 + zero
@ -183,7 +183,7 @@ class Frame(object):
(zero, sa, ca * v),
))
domega = _mxv(solutn, rates[::-1])
domega = mxv(solutn, rates[::-1])
drdtrt = array((
(zero, domega[0], domega[2]),
@ -191,11 +191,11 @@ class Frame(object):
(-domega[2], -domega[1], zero),
))
dRdt = _mxm(drdtrt, R)
dRdt = mxm(drdtrt, R)
if self._matrix is not None:
R = _mxm(self._matrix, R)
dRdt = _mxm(self._matrix, dRdt)
R = mxm(self._matrix, R)
dRdt = mxm(self._matrix, dRdt)
return R, dRdt
@ -215,7 +215,7 @@ class PlanetTopos(VectorFunction):
@classmethod
def from_latlon_distance(cls, frame, latitude, longitude, distance):
r = array((distance.au, 0.0, 0.0))
r = _mxv(rot_z(longitude.radians), _mxv(rot_y(-latitude.radians), r))
r = mxv(rot_z(longitude.radians), mxv(rot_y(-latitude.radians), r))
self = cls(frame, r)
self.latitude = latitude
@ -224,13 +224,13 @@ class PlanetTopos(VectorFunction):
def _at(self, t):
R, dRdt = self._frame.rotation_and_rate_at(t)
r = _mxv(_T(R), self._position_au)
v = _mxv(_T(dRdt), self._position_au) * DAY_S
r = mxv(_T(R), self._position_au)
v = mxv(_T(dRdt), self._position_au) * DAY_S
return r, v, None, None
def _snag_observer_data(self, observer_data, t):
R = self._frame.rotation_at(t)
observer_data.altaz_rotation = _mxmxm(
observer_data.altaz_rotation = mxmxm(
# TODO: Figure out how to produce this rotation directly
# from _position_au, to support situations where we were not
# given a latitude and longitude. If that is not feasible,

View File

@ -7,7 +7,7 @@ from .data.spice import inertial_frames
from .earthlib import compute_limb_angle, refract, reverse_terra
from .geometry import intersect_line_and_sphere
from .functions import (
_mxv, _mxm, _to_array, angle_between, from_spherical,
mxv, mxm, _to_array, angle_between, from_spherical,
length_of, rot_x, rot_z, to_spherical,
)
from .relativity import add_aberration, add_deflection
@ -290,7 +290,7 @@ class ICRF(object):
"""
if epoch is None:
vector = _mxv(_ECLIPJ2000, self.position.au)
vector = mxv(_ECLIPJ2000, self.position.au)
return Distance(vector)
position_au = self.position.au
@ -308,8 +308,8 @@ class ICRF(object):
_, d_eps = epoch._nutation_angles_radians
true_obliquity = epoch._mean_obliquity_radians + d_eps
rotation = _mxm(rot_x(- true_obliquity), epoch.M)
position_au = _mxv(rotation, position_au)
rotation = mxm(rot_x(- true_obliquity), epoch.M)
position_au = mxv(rotation, position_au)
return Distance(position_au)
def ecliptic_velocity(self):
@ -347,12 +347,12 @@ class ICRF(object):
def frame_xyz(self, frame):
"""Express this position as an (x,y,z) vector in a particular frame."""
R = frame.rotation_at(self.t)
return Distance(au=_mxv(R, self.position.au))
return Distance(au=mxv(R, self.position.au))
def frame_latlon(self, frame):
"""Return as longitude, latitude, and distance in the given frame."""
R = frame.rotation_at(self.t)
vector = _mxv(R, self.position.au)
vector = mxv(R, self.position.au)
d, lat, lon = to_spherical(vector)
return (Angle(radians=lat, signed=True),
Angle(radians=lon),

View File

@ -7,7 +7,7 @@ from numpy import (
from sgp4.api import SGP4_ERRORS, Satrec
from .constants import AU_KM, DAY_S, T0, tau
from .functions import _mxv, rot_x, rot_y, rot_z
from .functions import mxv, rot_x, rot_y, rot_z
from .io_timescale import _build_builtin_timescale
from .positionlib import ITRF_to_GCRS2
from .searchlib import _find_discrete, find_maxima
@ -313,8 +313,8 @@ def TEME_to_ITRF(jd_ut1, rTEME, vTEME, xp=0.0, yp=0.0):
rPEF = (R).dot(rTEME)
vPEF = (R).dot(vTEME) + _cross(angular_velocity, rPEF)
else:
rPEF = _mxv(R, rTEME)
vPEF = _mxv(R, vTEME) + _cross(angular_velocity, rPEF)
rPEF = mxv(R, rTEME)
vPEF = mxv(R, vTEME) + _cross(angular_velocity, rPEF)
if xp == 0.0 and yp == 0.0:
rITRF = rPEF

View File

@ -9,7 +9,7 @@ from .constants import ASEC2RAD, B1950, DAY_S, T0, tau
from .descriptorlib import reify
from .earthlib import sidereal_time, earth_rotation_angle
from .framelib import ICRS_to_J2000 as B
from .functions import _mxm, _mxmxm, _to_array, load_bundled_npy, rot_z
from .functions import mxm, mxmxm, _to_array, load_bundled_npy, rot_z
from .nutationlib import (
build_nutation_matrix, equation_of_the_equinoxes_complimentary_terms,
iau2000a_radians, mean_obliquity,
@ -599,7 +599,7 @@ class Time(object):
if N is None:
N = self.nutation_matrix()
return _mxmxm(N, P, B)
return mxmxm(N, P, B)
@reify
def MT(self):
@ -611,7 +611,7 @@ class Time(object):
# Calculate the Equation of Origins in cycles
eq_origins = (earth_rotation_angle(self.ut1) - self.gast / 24.0)
R = rot_z(2 * pi * eq_origins)
return _mxm(R, self.M)
return mxm(R, self.M)
@reify
def CT(self):

View File

@ -2,7 +2,7 @@ from numpy import einsum
from .constants import ASEC2RAD, tau
from .earthlib import terra
from .functions import _mxmxm, rot_x, rot_y, rot_z
from .functions import mxmxm, rot_x, rot_y, rot_z
from .units import Distance, Angle, _interpret_ltude
from .vectorlib import VectorFunction
@ -68,7 +68,7 @@ class Topos(VectorFunction):
def _altaz_rotation(self, t):
"""Compute the rotation from the ICRF into the alt-az system."""
R_lon = rot_z(- self.longitude.radians - t.gast * tau / 24.0)
return _mxmxm(self.R_lat, R_lon, t.M)
return mxmxm(self.R_lat, R_lon, t.M)
def _at(self, t):
"""Compute the GCRS position and velocity of this Topos at time `t`."""