Source code for spacepy.coordinates

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Implementation of Coords class functions for coordinate transformations

Authors: Josef Koller and Steven Morley
Institution: Los ALamos National Laboratory
Contact: smorley@lanl.gov

Copyright 2010 Los Alamos National Security, LLC.
"""

import numpy as np
from spacepy import help
import spacepy
import spacepy.time

__contact__ = 'Steven Morley, smorley@lanl.gov'

# -----------------------------------------------
# space coordinate class
# -----------------------------------------------
[docs]class Coords(object): ''' a = Coords( data, dtype, carsph, [units, ticks] ) A class holding spatial coordinates in Cartesian/spherical in units of Re and degrees .. note:: Although other units may be specified and will be carried through, most functions throughout SpacePy assume distances in Re and angles in degrees, regardless of specified units. Coordinate transforms are based on the IRBEM library; `its manual <http://svn.code.sf.net/p/irbem/code/trunk/manual/user_guide.html>`_ may prove useful. For a good reference on heliospheric and magnetospheric coordinate systems, see Franz & Harper, "Heliospheric Coordinate Systems", Planet. Space Sci., 50, pp 217-233, 2002 (https://doi.org/10.1016/S0032-0633(01)00119-2). Parameters ========== data : list or ndarray, dim = (n,3) coordinate points [X,Y,Z] or [rad, lat, lon] dtype : string coordinate system; possible values are: * **GDZ** (Geodetic; WGS84), * **GEO** (Geographic Coordinate System), * **GSM** (Geocentric Solar Magnetospheric), * **GSE** (Geocentric Solar Ecliptic), * **SM** (Solar Magnetic), * **GEI** (Geocentric Equatorial Inertial; True-of-Date), * **MAG** (Geomagnetic Coordinate System), * **SPH** (Spherical Coordinate System), * **RLL** (Radius, Latitude, Longitude; Geodetic) carsph : string Cartesian or spherical, 'car' or 'sph' units : list of strings, optional standard are ['Re', 'Re', 'Re'] or ['Re', 'deg', 'deg'] depending on the carsph content. See note. ticks : Ticktock instance, optional used for coordinate transformations (see a.convert) Returns ======= out : Coords instance instance with a.data, a.carsph, etc. See Also ======== spacepy.time.Ticktock Examples ======== >>> from spacepy import coordinates as coord >>> cvals = coord.Coords([[1,2,4],[1,2,2]], 'GEO', 'car') >>> cvals.x # returns all x coordinates array([1, 1]) >>> from spacepy.time import Ticktock >>> cvals.ticks = Ticktock(['2002-02-02T12:00:00', '2002-02-02T12:00:00'], 'ISO') # add ticks >>> newcoord = cvals.convert('GSM', 'sph') >>> newcoord .. currentmodule:: spacepy.coordinates .. autosummary:: ~Coords.append ~Coords.convert ~Coords.from_skycoord ~Coord.to_skycoord .. automethod:: append .. automethod:: convert .. automethod:: from_skycoord .. automethod:: to_skycoord ''' Re = 6371200.0 # meters def __init__(self, data, dtype, carsph, units=None, ticks=None): from . import irbempy as op from spacepy.irbempy import SYSAXES_TYPES as typedict if isinstance(data[0], (float, int)): self.data = np.array([data]) else: self.data = np.array(data) dtype = dtype.upper() carsph = carsph.lower() assert dtype in list(typedict.keys()), 'This dtype='+dtype+' is not supported. Only '+str(list(typedict.keys())) assert carsph in ['car','sph'], 'This carsph='+str(carsph)+' is not supported. Only "car" or "sph"' onerawarn = """Coordinate conversion to an ONERA-compatible system is required for any ONERA calls.""" # add ticks if ticks: assert len(ticks) == len(self.data), 'Ticktock dimensions seem off' self.ticks = ticks # GEO,sph and SPH,sph are the same if dtype == 'GEO' and carsph == 'sph': dtype = 'SPH' self.sysaxes = typedict[dtype][carsph] #if self.sysaxes >= 10 and self.sysaxes < 20: #need sph2car # try: # self.data = op.sph2car(self.data) # self.sysaxes -= 10 # except: # print onerawarn # self.sysaxes = None #if self.sysaxes >= 20: #need car2sph # try: # self.data = op.car2sph(self.data) # self.sysaxes -= 20 # except: # print onerawarn # self.sysaxes = None self.dtype = dtype self.carsph = carsph # setup units if units is None and carsph == 'car': # use standard units self.units = ['Re', 'Re', 'Re'] elif units is None and carsph == 'sph': self.units = ['Re', 'deg', 'deg'] else: self.units = units if dtype == 'GDZ' and carsph == 'sph': self.units = ['km', 'deg', 'deg'] # setup x,y,z etc if carsph == 'car': self.x = spacepy.datamodel.dmcopy(self.data[:, 0]) self.y = spacepy.datamodel.dmcopy(self.data[:, 1]) self.z = spacepy.datamodel.dmcopy(self.data[:, 2]) else: self.radi = spacepy.datamodel.dmcopy(self.data[:, 0]) self.lati = spacepy.datamodel.dmcopy(self.data[:, 1]) self.long = spacepy.datamodel.dmcopy(self.data[:, 2]) ## setup list for onera compatibility #self.sysaxes = op.get_sysaxes(dtype, carsph) self.shape = np.shape(self.data) return None # ----------------------------------------------- def __str__(self): ''' a.__str__() or a Will be called when printing Coords instance a Returns ======== out : string string represenation of the instance Examples ======== >>> from spacepy.coordinates import Coords >>> y = Coords([[1,2,4],[1,2,2]], 'GEO', 'car') >>> y Coords( [[1 2 4] [1 2 2]] ), dtype=GEO,car, units=['Re', 'Re', 'Re'] ''' rstr = "Coords( {0} , '{1}', '{2}')".format(self.data.tolist(), self.dtype, self.carsph) return rstr __repr__ = __str__ # ----------------------------------------------- def __getitem__(self, idx): ''' a.__getitem__(idx) or a[idx] Will be called when requesting items in this instance Parameters ========== idx : int integer numbers as index Returns ======= out : numpy array new values Examples ======== >>> from spacepy.coordinates import Coords >>> y = Coords([[1,2,4],[1,2,2]], 'GEO', 'car') >>> y[0] Coords( [[1 2 4]] ), dtype=GEO,car, units=['Re', 'Re', 'Re'] ''' arr = np.array(self.data) t_select = self.ticks[idx] if self.ticks else self.ticks return Coords(arr[idx].tolist(), self.dtype, self.carsph, self.units, t_select) # ----------------------------------------------- def __setitem__(self, idx, vals): ''' a.__setitem__(idx, vals) or a[idx] = vals Will be called setting items in this instance Parameters ========== idx : int integer numbers as index vals : numpy array or list new values Examples ======== >>> from spacepy.coordinates import Coords >>> y = Coords([[1,2,4],[1,2,2]], 'GEO', 'car') >>> y[1] = [9,9,9] >>> y Coords( [[1 2 4] [9 9 9]] ), dtype=GEO,car, units=['Re', 'Re', 'Re'] ''' self.data[idx] = vals return # ----------------------------------------------- def __len__(self): ''' a.__len__() or len(a) Will be called when requesting the length, i.e. number of items Returns ======== out : int length Examples ======== >>> from spacepy.coordinates import Coords >>> y = Coords([[1,2,4],[1,2,2]], 'GEO', 'car') >>> len(y) 2 ''' if isinstance(self.data, (list, np.ndarray)): return len(self.data) else: return 1 return # -----------------------------------------------
[docs] def convert(self, returntype, returncarsph): '''Create a new Coords instance with new coordinate types Parameters ========== returntype : string coordinate system, possible are GDZ, GEO, GSM, GSE, SM, GEI, MAG, SPH, RLL returncarsph : string coordinate type, possible 'car' for Cartesian and 'sph' for spherical Returns ======= out : Coords object Coords object in the new coordinate system Examples ======== >>> from spacepy.coordinates import Coords >>> y = Coords([[1,2,4],[1,2,2]], 'GEO', 'car') >>> from spacepy.time import Ticktock >>> y.ticks = Ticktock(['2002-02-02T12:00:00', '2002-02-02T12:00:00'], 'ISO') >>> x = y.convert('SM','car') >>> x Coords( [[ 0.81134097 2.6493305 3.6500375 ] [ 0.92060408 2.30678864 1.68262126]] ), dtype=SM,car, units=['Re', 'Re', 'Re'] ''' from . import irbempy as op from spacepy.irbempy import SYSAXES_TYPES as typedict #check return type/system is supported #if not (typedict[returntype][returncarsph]): # raise NotImplementedError('System {0} is not supported in {1} coordinates'.format(returntype, returncarsph)) # no change necessary if (self.dtype == returntype) and (self.carsph == returncarsph): return self # only car2sph or sph2car is needed if (self.dtype == returntype) and (self.carsph != returncarsph): if returncarsph == 'car': carsph = 'car' units = [self.units[0]]*3 data = op.sph2car(self.data) else: carsph = 'sph' units = [self.units[0], 'deg','deg'] data = op.car2sph(self.data) return Coords(data, self.dtype, carsph, units, self.ticks) # check the length of ticks and do the more complex conversions if self.ticks: assert len(self.ticks) == len(self), 'Ticktock dimension does not match Coords dimensions' # check if car2sph is needed first for oneralib compatibility if (self.sysaxes is None) : # a car2sph or sph2car is needed if self.carsph == 'sph': carsph = 'car' units = [self.units[0]]*3 data = op.sph2car(self.data) else: carsph = 'sph' units = [self.units[0], 'deg','deg'] data = op.car2sph(self.data) else: data = self.data units = self.units carsph = self.carsph NewCoords = Coords(data, self.dtype, carsph, units, self.ticks) # now convert to other coordinate system if (self.dtype != returntype) : assert NewCoords.ticks, "Time information required; add a.ticks attribute" NewCoords.data = op.coord_trans( NewCoords, returntype, returncarsph) NewCoords.dtype = returntype NewCoords.carsph = returncarsph NewCoords.sysaxes = op.get_sysaxes(returntype, returncarsph) # fix corresponding attributes if returncarsph == 'sph': NewCoords.units = [units[0], 'deg','deg'] for k in ('x', 'y', 'z'): if hasattr(NewCoords, k): delattr(NewCoords, k) NewCoords.radi = NewCoords.data[:,0] NewCoords.lati = NewCoords.data[:,1] NewCoords.long = NewCoords.data[:,2] else: # 'car' NewCoords.units = [units[0]]*3 for k in ('radi', 'lati', 'long'): if hasattr(NewCoords, k): delattr(NewCoords, k) NewCoords.x = NewCoords.data[:,0] NewCoords.y = NewCoords.data[:,1] NewCoords.z = NewCoords.data[:,2] return NewCoords
# ----------------------------------------------- def __getstate__(self): ''' Is called when pickling See Also http://docs.python.org/library/pickle.html ''' odict = self.__dict__.copy() # copy the dict since we change it return odict def __setstate__(self, dict): ''' Is called when unpickling See Also http://docs.python.org/library/pickle.html ''' self.__dict__.update(dict) return # -----------------------------------------------
[docs] def append(self, other): '''Append another Coords instance to the current one Parameters ========== other : Coords instance Coords instance to append Examples ======== ''' data = list(self.data) otherdata = other.convert(self.dtype, self.carsph) data.extend(list(otherdata.data)) newobj = Coords(data, dtype=self.dtype, carsph=self.carsph) return newobj
# -----------------------------------------------
[docs] def to_skycoord(self): ''' Create an Astropy SkyCoord instance based on this instance Returns ======= out : `astropy.coordinates.SkyCoord` This coordinate as an Astropy SkyCoord Notes ===== This method requires Astropy to be installed. This method uses the GEO coordinate frame as the common frame between the two libraries. ''' if self.ticks is None: raise ValueError("This method requires the attribute `ticks` to be specified.") try: import astropy.coordinates import astropy.time except ImportError as e: raise e.__class__("This method requires Astropy to be installed.") coord = self.convert('GEO', 'car') # GEO is Astropy ITRS data = coord.data * self.Re # convert to meters obstime = coord.ticks.APT return astropy.coordinates.SkyCoord(x=data[:, 0], y=data[:, 1], z=data[:, 2], unit='m', frame='itrs', obstime=obstime)
# -----------------------------------------------
[docs] @classmethod def from_skycoord(cls, skycoord): ''' Create a Coords instance from an Astropy SkyCoord instance Parameters ========== skycoord : `astropy.coordinates.SkyCoord` The coordinate to be converted Returns ======= out : Coords instance The converted coordinate Notes ===== This method requires Astropy to be installed. This method uses the GEO coordinate frame as the common frame between the two libraries. ''' try: import astropy.coordinates except ImportError as e: raise e.__class__("This method requires Astropy to be installed.") skycoord = astropy.coordinates.SkyCoord(skycoord).itrs # Astropy ITRS is GEO data = (skycoord.cartesian.xyz.to('m').value / cls.Re).T # convert Cartesian to Re units ticks = spacepy.time.Ticktock(skycoord.obstime) return cls(data, 'GEO', 'car', ticks=ticks)
[docs]def quaternionNormalize(Qin, scalarPos='last'): ''' Given an input quaternion (or array of quaternions), return the unit quaternion Parameters ========== vec : array_like input quaternion to normalize Returns ======= out : array_like normalized quaternion Examples ======== >>> import spacepy.coordinates >>> spacepy.coordinates.quaternionNormalize([0.707, 0, 0.707, 0.2]) array([ 0.69337122, 0. , 0.69337122, 0.19614462]) ''' w = {'first': 0, 'last': 3}.get(scalarPos.lower()) if w is None: raise NotImplementedError('quaternionNormalize: scalarPos must be set to "First" or "Last"') Quse = np.asanyarray(Qin).astype(float) squeeze = Quse.ndim < 2 if squeeze: Quse = np.expand_dims(Quse, axis=0) magn = np.linalg.norm(Quse, axis=-1) # Find places where the magnitude is tiny, convert to unit real idx = np.where(magn <= 1e-12) magn[idx] = 1 outarr = Quse / np.expand_dims(magn, axis=magn.ndim) outarr[idx + (None,)] = 0 outarr[idx + (w,)] = 1 if squeeze: outarr = outarr[0, ...] return outarr
[docs]def quaternionRotateVector(Qin, Vin, scalarPos='last', normalize=True): ''' Given quaternions and vectors, return the vectors rotated by the quaternions Parameters ========== Qin : array_like input quaternion to rotate by Vin : array-like input vector to rotate Returns ======= out : array_like rotated vector Examples ======== >>> import spacepy.coordinates >>> import numpy as np >>> vec = [1, 0, 0] >>> quat_wijk = [np.sin(np.pi/4), 0, np.sin(np.pi/4), 0.0] >>> quat_ijkw = [0.0, np.sin(np.pi/4), 0, np.sin(np.pi/4)] >>> spacepy.coordinates.quaternionRotateVector(quat_ijkw, vec) array([ 0., 0., -1.]) >>> spacepy.coordinates.quaternionRotateVector( ... quat_wijk, vec, scalarPos='first') array([ 0., 0., -1.]) See Also ======== quaternionMultiply ''' if scalarPos.lower()=='last': i, j, k = 0, 1, 2 w = 3 elif scalarPos.lower()=='first': i, j, k = 1, 2, 3 w = 0 else: raise NotImplementedError('quaternionRotateVector: scalarPos must be set to "First" or "Last"') Quse = np.asanyarray(Qin).astype(float) if normalize: Quse = quaternionNormalize(Quse, scalarPos=scalarPos) Vuse = np.asanyarray(Vin).astype(float) try: Quse.shape[1] except IndexError: Quse = np.asanyarray([Quse]) try: Vuse.shape[1] except IndexError: Vuse = np.asanyarray([Vuse]) try: assert Vuse.shape[0]==Quse.shape[0] except AssertionError: raise ValueError('quaternionRotateVector: Input vector array must have same length as input quaternion array') outarr = np.empty((Quse.shape[0], 3)) for idx, row in enumerate(Quse): Vrow = Vuse[idx] ii, jj, kk, ww = row[i]**2, row[j]**2, row[k]**2, row[w]**2 iw, jw, kw = row[i]*row[w], row[j]*row[w], row[k]*row[w] ij, ik, jk = row[i]*row[j], row[i]*row[k], row[j]*row[k] outarr[idx, 0] = ww*Vrow[0] + 2*jw*Vrow[2] - 2*kw*Vrow[1] + ii*Vrow[0] + 2*ij*Vrow[1] + 2*ik*Vrow[2] - kk*Vrow[0] - jj*Vrow[0] outarr[idx, 1] = 2*ij*Vrow[0] + jj*Vrow[1] + 2*jk*Vrow[2] + 2*kw*Vrow[0] - kk*Vrow[1] + ww*Vrow[1] - 2*iw*Vrow[2] - ii*Vrow[1] outarr[idx, 2] = 2*ik*Vrow[0] + 2*jk*Vrow[1] + kk*Vrow[2] - 2.0*jw*Vrow[0] - jj*Vrow[2] + 2.0*iw*Vrow[1] - ii*Vrow[2] + ww*Vrow[2] return outarr.squeeze()
[docs]def quaternionMultiply(Qin1, Qin2, scalarPos='last'): ''' Given quaternions, return the product, i.e. Qin1*Qin2 Parameters ========== Qin1 : array_like input quaternion, first position Qin2 : array-like input quaternion, second position Returns ======= out : array_like quaternion product Examples ======== >>> import spacepy.coordinates >>> import numpy as np >>> vecX = [1, 0, 0] #shared X-axis >>> vecZ = [0, 0, 1] #unshared, but similar, Z-axis >>> quat_eci_to_gsm = [-0.05395384, 0.07589845, -0.15172533, 0.98402634] >>> quat_eci_to_gse = [ 0.20016056, 0.03445775, -0.16611386, 0.96496352] >>> quat_gsm_to_eci = spacepy.coordinates.quaternionConjugate( ... quat_eci_to_gsm) >>> quat_gse_to_gsm = spacepy.coordinates.quaternionMultiply( ... quat_gsm_to_eci, quat_eci_to_gse) >>> spacepy.coordinates.quaternionRotateVector(quat_gse_to_gsm, vecX) array([ 1.00000000e+00, 1.06536725e-09, -1.16892107e-08]) >>> spacepy.coordinates.quaternionRotateVector(quat_gse_to_gsm, vecZ) array([ 1.06802834e-08, -4.95669027e-01, 8.68511494e-01]) ''' if scalarPos.lower()=='last': i, j, k = 0, 1, 2 w = 3 elif scalarPos.lower()=='first': i, j, k = 1, 2, 3 w = 0 else: raise NotImplementedError('quaternionMultiply: scalarPos must be set to "First" or "Last"') Quse1 = np.asanyarray(Qin1).astype(float) Quse2 = np.asanyarray(Qin2).astype(float) try: Quse1.shape[1] except IndexError: Quse1 = np.asanyarray([Quse1]) try: Quse2.shape[1] except IndexError: Quse2 = np.asanyarray([Quse2]) try: assert Quse2.shape==Quse1.shape except AssertionError: raise ValueError('quaternionMultiply: Input quaternion arrays must have same length') outarr = np.empty_like(Quse1) for idx, row in enumerate(Quse1): row1 = row row2 = Quse2[idx] #vector components outarr[idx, i] = row1[w]*row2[i] + row1[i]*row2[w] + row1[j]*row2[k] - row1[k]*row2[j] outarr[idx, j] = row1[w]*row2[j] - row1[i]*row2[k] + row1[j]*row2[w] + row1[k]*row2[i] outarr[idx, k] = row1[w]*row2[k] + row1[i]*row2[j] - row1[j]*row2[i] + row1[k]*row2[w] #real part outarr[idx, w] = row1[w]*row2[w] - row1[i]*row2[i] - row1[j]*row2[j] - row1[k]*row2[k] return outarr.squeeze()
[docs]def quaternionConjugate(Qin, scalarPos='last'): ''' Given an input quaternion (or array of quaternions), return the conjugate Parameters ========== Qin : array_like input quaternion to conjugate Returns ======= out : array_like conjugate quaternion Examples ======== >>> import spacepy.coordinates >>> spacepy.coordinates.quaternionConjugate( ... [0.707, 0, 0.707, 0.2], scalarPos='last') array([-0.707, -0. , -0.707, 0.2 ]) See Also ======== quaternionMultiply ''' if scalarPos.lower()=='last': i, j, k = 0, 1, 2 w = 3 elif scalarPos.lower()=='first': i, j, k = 1, 2, 3 w = 0 else: raise NotImplementedError('quaternionConjugate: scalarPos must be set to "First" or "Last"') Quse = np.asanyarray(Qin).astype(float) try: Quse.shape[1] except IndexError: Quse = np.asanyarray([Quse]) outarr = np.empty_like(Quse) for idx, row in enumerate(Quse): outarr[idx, i] = -row[i] outarr[idx, j] = -row[j] outarr[idx, k] = -row[k] outarr[idx, w] = row[w] return outarr.squeeze()
[docs]def quaternionFromMatrix(matrix, scalarPos='last'): ''' Given an input rotation matrix, return the equivalent quaternion The output has one fewer axis than the input (the last axis) and the shape is otherwise unchanged, allowing multi-dimensional matrix input. Parameters ========== matrix : array_like input rotation matrix or array of matrices Returns ======= out : array_like Quaternions representing the same rotation as the input rotation matrices. Other Parameters ================ scalarPos : str Location of the scalar component of the output quaternion, either 'last' (default) or 'first'. Raises ====== NotImplementedError for invalid values of ``scalarPos`` ValueError for inputs which are obviously not valid 3D rotation matrices or arrays thereof: if the size doesn't end in (3, 3), if the matrix is not orthogonal, or not a proper rotation. See Also ======== quaternionToMatrix Notes ===== .. versionadded:: 0.2.2 No attempt is made to resolve the sign ambiguity; in particular, conversions of very similar matrices may result in equivalent quaternions with the opposite sign. This may have implications for interpolating a sequence of quaternions. The conversion of a rotation matrix to a quaternion suffers from some of the same disadvantages inherent to rotation matrices, such as potential numerical instabilities. Working in quaternion space as much as possible is recommended. There are several algorithms; the most well-known algorithm for this conversion is Shepperd's [#Shepperd]_, although the many "rediscoveries" indicate it is not sufficiently well-known. This function uses the method of Bar-Itzhack [#BarItzhack]_ (version 3), which should be resistant to small errors in the rotation matrix. As a result, the input checking is quite coarse and will likely accept many matrices that do not represent valid rotations. Also potentially of interest, although not implemented here, is Sarabandi and Thomas [#Sarabandi]_. References ========== .. [#Shepperd] S.W. Shepperd, "Quaternion from rotation matrix," Journal of Guidance and Control, Vol. 1, No. 3, pp. 223-224, 1978, `doi:10.2514/3.55767b <https://doi.org/10.2514/3.55767b>`_ .. [#BarItzhack] I. Y. Bar-Itzhack, "New method for extracting the quaternion from a rotation matrix", AIAA Journal of Guidance, Control and Dynamics, 23 (6): 1085–1087, `doi:10.2514/2.4654 <https://doi.org/10.2514/2.4654>`_ .. [#Sarabandi] S. Sarabandi and F. Thomas, "Accurate Computation of Quaternions from Rotation Matrices", In: Lenarcic J., Parenti-Castelli V. (eds) Advances in Robot Kinematics 2018, Springer. `doi:10.1007/978-3-319-93188-3_5 <https://doi.org/10.1007/978-3-319-93188-3_5>`_ Examples ======== >>> import spacepy.coordinates >>> spacepy.coordinates.quaternionFromMatrix( ... [[ 0., 0., 1.], ... [ 1., 0., 0.], ... [ 0., 1., 0.]]) array([0.5, 0.5, 0.5, 0.5]) ''' if scalarPos.lower() not in ('last', 'first'): raise NotImplementedError( 'quaternionFromMatrix: scalarPos must be set to "First" or "Last"') matrix = np.asanyarray(matrix) if len(matrix.shape) < 2 or matrix.shape[-2:] != (3, 3): raise ValueError( 'Input does not appear to be 3D rotation matrix, wrong size.') for i in np.ndindex(matrix.shape[:-2]): m = matrix[i + (Ellipsis,)] if not np.allclose(np.dot(m, m.transpose()), np.identity(3), atol=0.2): raise ValueError('Input rotation matrix{} not orthogonal.'.format( '' if len(matrix.shape) == 2 else ' at ' + str(i))) det = np.linalg.det(m) if det < 0: raise ValueError('Input rotation matrix at {} not proper.' .format(str(i))) inshape = matrix.shape # Flatten out most dimensions, easier to index matrix = np.reshape(matrix, (-1, 3, 3)) # Indexing in Bar-Itzhack is reversed relative to numpy k = ( matrix[..., 0, 0] - matrix[..., 1, 1] - matrix[..., 2, 2], matrix[..., 0, 1] + matrix[..., 1, 0], matrix[..., 0, 2] + matrix[..., 2, 0], matrix[..., 2, 1] - matrix[..., 1, 2], # row 0 matrix[..., 0, 1] + matrix[..., 1, 0], matrix[..., 1, 1] - matrix[..., 0, 0] - matrix[..., 2, 2], matrix[..., 1, 2] + matrix[..., 2, 1], matrix[..., 0, 2] - matrix[..., 2, 0], # row 1 matrix[..., 0, 2] + matrix[..., 2, 0], matrix[..., 1, 2] + matrix[..., 2, 1], matrix[..., 2, 2] - matrix[..., 0, 0] - matrix[..., 1, 1], matrix[..., 1, 0] - matrix[..., 0, 1], #row 2 matrix[..., 2, 1] - matrix[..., 1, 2], matrix[..., 0, 2] - matrix[..., 2, 0], matrix[..., 1, 0] - matrix[..., 0, 1], matrix[..., 0, 0] + matrix[..., 1, 1] + matrix[..., 2, 2] #row 3 ) # Stack together on a new axis after the input sample, then # split that axis into the 4x4 k arrays k = np.reshape(np.stack(k, axis=1), (-1, 4, 4)) k = k / 3. evals, evects = np.linalg.eig(k) evects = np.real(evects) idx = np.argmax(evals, axis=-1) Qout = evects[np.mgrid[0:len(idx)], :, idx] if scalarPos.lower() == 'first': Qout = np.roll(Qout, 1, axis=-1) # Recover original shape Qout = np.reshape(Qout, inshape[:-2] + (4,)) return Qout
[docs]def quaternionToMatrix(Qin, scalarPos='last', normalize=True): ''' Given an input quaternion, return the equivalent rotation matrix. The output has one more axis than the input (the last axis) and the shape is otherwise unchanged, allowing multi-dimensional quaternion input. Parameters ========== Qin : array_like input quaternion or array of quaternions, must be normalized. Returns ======= out : array_like Rotation matrix Other Parameters ================ scalarPos : str Location of the scalar component of the input quaternion, either 'last' (default) or 'first'. normalize : True Normalize input quaternions before conversion (default). If False, raises error for non-normalized. Raises ====== NotImplementedError for invalid values of ``scalarPos``. ValueError for inputs which are not valid normalized quaternions or arrays thereof: if the size doesn't end in (4), if the quaternion is not normalized and ``normalize`` is False. See Also ======== quaternionFromMatrix Notes ===== .. versionadded:: 0.2.2 Implementation of the Euler–Rodrigues formula. Examples ======== >>> import spacepy.coordinates >>> spacepy.coordinates.quaternionToMatrix([0.5, 0.5, 0.5, 0.5]) array([[ 0., 0., 1.], [ 1., 0., 0.], [ 0., 1., 0.]]) ''' if scalarPos.lower() not in ('last', 'first'): raise NotImplementedError( 'quaternionToMatrix: scalarPos must be set to "First" or "Last"') Qin = np.asanyarray(Qin) if Qin.shape[-1] != 4: raise ValueError('Input does not appear to be quaternion, wrong size.') if normalize: Qin = quaternionNormalize(Qin, scalarPos=scalarPos) if scalarPos.lower() == 'first': Qin = np.roll(Qin, -1, axis=-1) if not np.allclose(np.sum(Qin ** 2, axis=-1), 1): raise ValueError('Input quaternion not normalized.') # Maintain dimensions at end for stacking a, b, c, d = Qin[..., -1:], Qin[..., 0:1], Qin[..., 1:2], Qin[..., 2:3] # Output array, "flattened" # This can probably be written with a clever vector notation... out = np.concatenate(( a ** 2 + b ** 2 - c ** 2 - d ** 2, 2 * (b * c - a * d), 2 * (b * d + a * c), 2 * (b * c + a * d), a ** 2 + c ** 2 - b ** 2 - d ** 2, 2 * (c * d - a * b), 2 * (b * d - a * c), 2 * (c * d + a * b), a ** 2 + d ** 2 - b ** 2 - c ** 2, ), axis=-1) # And last two dims are 3x3 array for matrix out = np.reshape(out, out.shape[:-1] + (3, 3)) return out