Source code for pycraf.geometry.geometry

#!/usr/bin/python
# -*- coding: utf-8 -*-

from __future__ import (
    absolute_import, unicode_literals, division, print_function
    )

from functools import reduce
from astropy import units as apu
import numpy as np
from .cygeometry import (
    true_angular_distance_cython, great_circle_bearing_cython
    )
from .. import utils


__all__ = [
    'true_angular_distance', 'great_circle_bearing',
    'cart_to_sphere', 'sphere_to_cart',
    'Rx', 'Ry', 'Rz', 'multiply_matrices',
    'rotaxis_from_rotmat', 'rotmat_from_rotaxis', 'eulerangle_from_rotmat',
    ]


[docs]@utils.ranged_quantity_input( l1=(None, None, apu.deg), b1=(-90, 90, apu.deg), l2=(None, None, apu.deg), b2=(-90, 90, apu.deg), strip_input_units=True, output_unit=apu.deg, ) def true_angular_distance(l1, b1, l2, b2): ''' True angular distance between points (l1, b1) and (l2, b2). Based on Vincenty formula (http://en.wikipedia.org/wiki/Great-circle_distance). This was spotted in astropy source code. Parameters ---------- l1, b1 : `~astropy.units.Quantity` Longitude/Latitude of point 1 [deg] l2, b2 : `~astropy.units.Quantity` Longitude/Latitude of point 2 [deg] Returns ------- adist : `~astropy.units.Quantity` True angular distance [deg] ''' # sin_diff_lon = np.sin(np.radians(l2 - l1)) # cos_diff_lon = np.cos(np.radians(l2 - l1)) # sin_lat1 = np.sin(np.radians(b1)) # sin_lat2 = np.sin(np.radians(b2)) # cos_lat1 = np.cos(np.radians(b1)) # cos_lat2 = np.cos(np.radians(b2)) # num1 = cos_lat2 * sin_diff_lon # num2 = cos_lat1 * sin_lat2 - sin_lat1 * cos_lat2 * cos_diff_lon # denominator = sin_lat1 * sin_lat2 + cos_lat1 * cos_lat2 * cos_diff_lon # return np.degrees(np.arctan2( # np.sqrt(num1 ** 2 + num2 ** 2), denominator # )) return true_angular_distance_cython(l1, b1, l2, b2)
[docs]@utils.ranged_quantity_input( l1=(None, None, apu.deg), b1=(-90, 90, apu.deg), l2=(None, None, apu.deg), b2=(-90, 90, apu.deg), strip_input_units=True, output_unit=apu.deg, ) def great_circle_bearing(l1, b1, l2, b2): ''' Great circle bearing between points (l1, b1) and (l2, b2). Parameters ---------- l1, b1 : `~astropy.units.Quantity` Longitude/Latitude of point 1 [deg] l2, b2 : `~astropy.units.Quantity` Longitude/Latitude of point 2 [deg] Returns ------- bearing : `~astropy.units.Quantity` Great circle bearing [deg] ''' # diff_lon_rad = np.radians(l2 - l1) # b1_rad, b2_rad = np.radians(b1), np.radians(b2) # sin_diff_lon = np.sin(diff_lon_rad) # cos_diff_lon = np.cos(diff_lon_rad) # sin_lat1 = np.sin(b1_rad) # sin_lat2 = np.sin(b2_rad) # cos_lat1 = np.cos(b1_rad) # cos_lat2 = np.cos(b2_rad) # a = cos_lat2 * sin_diff_lon # b = cos_lat1 * sin_lat2 - sin_lat1 * cos_lat2 * cos_diff_lon # return np.degrees(np.arctan2(a, b)) return great_circle_bearing_cython(l1, b1, l2, b2)
def _cart_to_sphere(x, y, z, broadcast_arrays=True): r = np.sqrt(x ** 2 + y ** 2 + z ** 2) theta = 90 - np.degrees(np.arccos(z / r)) phi = np.degrees(np.arctan2(y, x)) if broadcast_arrays: r, phi, theta = np.broadcast_arrays(r, phi, theta) return r, phi, theta
[docs]@utils.ranged_quantity_input( x=(None, None, apu.m), y=(None, None, apu.m), z=(None, None, apu.m), strip_input_units=True, output_unit=(apu.m, apu.deg, apu.deg) ) def cart_to_sphere(x, y, z, broadcast_arrays=True): ''' Spherical coordinates from Cartesian representation. Parameters ---------- x, y, z : `~astropy.units.Quantity` Cartesian position [m] Returns ------- r : `~astropy.units.Quantity` Radial distance [m] phi : `~astropy.units.Quantity` Azimuth [deg] theta : `~astropy.units.Quantity` Elevation [deg] broadcast_arrays : boolean, optional If 'True', output arrays will be broadcasted to dense matrices, otherwise the returned arrays will be a sparse representation (default: True) Notes ----- Unlike with the mathematical definition, `theta` is not the angle to the (positive) `z` axis, but the elevation above the `x`-`y` plane. ''' return _cart_to_sphere(x, y, z, broadcast_arrays=broadcast_arrays)
def _sphere_to_cart(r, phi, theta, broadcast_arrays=True): c_t, s_t = np.cos(np.radians(theta)), np.sin(np.radians(theta)) c_p, s_p = np.cos(np.radians(phi)), np.sin(np.radians(phi)) x = r * c_t * c_p y = r * c_t * s_p z = r * s_t if broadcast_arrays: x, y, z = np.broadcast_arrays(x, y, z) return x, y, z
[docs]@utils.ranged_quantity_input( r=(None, None, apu.m), theta=(-90, 90, apu.deg), phi=(None, None, apu.deg), strip_input_units=True, output_unit=(apu.m, apu.m, apu.m) ) def sphere_to_cart(r, phi, theta, broadcast_arrays=True): ''' Spherical coordinates from Cartesian representation. Parameters ---------- r : `~astropy.units.Quantity` Radial distance [m] phi : `~astropy.units.Quantity` Azimuth [deg] theta : `~astropy.units.Quantity` Elevation [deg] broadcast_arrays : boolean, optional If 'True', output arrays will be broadcasted to dense matrices, otherwise the returned arrays will be a sparse representation (default: True) Returns ------- x, y, z : `~astropy.units.Quantity` Cartesian position [m] Notes ----- Unlike with the mathematical definition, `theta` is not the angle to the (positive) `z` axis, but the elevation above the `x`-`y` plane. ''' return _sphere_to_cart(r, phi, theta, broadcast_arrays=broadcast_arrays)
def _Rx(angle): angle_rad = np.radians(angle) sin_a, cos_a = np.broadcast_arrays( np.sin(angle_rad.flat), np.cos(angle_rad.flat) ) o, z = np.ones_like(sin_a), np.zeros_like(sin_a) R = np.array([[o, z, z], [z, cos_a, -sin_a], [z, sin_a, cos_a]]) return R.swapaxes(0, 2).reshape(angle_rad.shape + (3, 3)) def _Ry(angle): angle_rad = np.radians(angle) sin_a, cos_a = np.broadcast_arrays( np.sin(angle_rad.flat), np.cos(angle_rad.flat) ) o, z = np.ones_like(sin_a), np.zeros_like(sin_a) R = np.array([[cos_a, z, sin_a], [z, o, z], [-sin_a, z, cos_a]]) return R.swapaxes(0, 2).reshape(angle_rad.shape + (3, 3)) def _Rz(angle): angle_rad = np.radians(angle) sin_a, cos_a = np.broadcast_arrays( np.sin(angle_rad.flat), np.cos(angle_rad.flat) ) o, z = np.ones_like(sin_a), np.zeros_like(sin_a) R = np.array([[cos_a, -sin_a, z], [sin_a, cos_a, z], [z, z, o]]) return R.swapaxes(0, 2).reshape(angle_rad.shape + (3, 3))
[docs]@utils.ranged_quantity_input( angle=(None, None, apu.deg), strip_input_units=True ) def Rx(angle): ''' Construct rotation matrix about x-axis. Parameters ---------- angle : `~astropy.units.Quantity` Rotation angle [deg] Returns ------- Rx : `~numpy.array` Rotation matrix [no units!] Notes ----- Broadcasting is supported. Rotation matrices can be (matrix-multiplied) with `~numpy.dot`, however, if you want to multiply stacks of matrices use the new `~numpy.matmul` function. ''' return _Rx(angle)
[docs]@utils.ranged_quantity_input( angle=(None, None, apu.deg), strip_input_units=True ) def Ry(angle): ''' Construct rotation matrix about y-axis. Parameters ---------- angle : `~astropy.units.Quantity` Rotation angle [deg] Returns ------- Ry : `~numpy.array` Rotation matrix [no units!] Notes ----- Broadcasting is supported. Rotation matrices can be (matrix-multiplied) with `~numpy.dot`, however, if you want to multiply stacks of matrices use the new `~numpy.matmul` function. ''' return _Ry(angle)
[docs]@utils.ranged_quantity_input( angle=(None, None, apu.deg), strip_input_units=True ) def Rz(angle): ''' Construct rotation matrix about z-axis. Parameters ---------- angle : `~astropy.units.Quantity` Rotation angle [deg] Returns ------- Rz : `~numpy.array` Rotation matrix [no units!] Notes ----- Broadcasting is supported. Rotation matrices can be (matrix-multiplied) with `~numpy.dot`, however, if you want to multiply stacks of matrices use the new `~numpy.matmul` function. ''' return _Rz(angle)
[docs]def multiply_matrices(*matrices): ''' Matrix-multiply the matrices in the given list. Parameters ---------- matrices : list of `~numpy.array` List of (stacks of) rotation matrices. The order of the multiplication is `R = R1.R2.R3`, if `*matrices = (R1, R2, R3)`. (This means, that `R3` is applied first.) Returns ------- R : `~numpy.array` Resulting rotation matrix [no units!] Notes ----- Broadcasting is supported, i.e., each element of `matrices` can be a stack of rotation matrices (that are of course not multiplied internally). ''' # Note, rotation matrices are associative, therefore we can replace # return reduce(lambda a, b: np.matmul(b, a), matrices[::-1]) # by return reduce(np.matmul, matrices)
def _rotmat_from_rotaxis(rotax_x, rotax_y, rotax_z, angle_deg): sh = np.broadcast(angle_deg, rotax_x, rotax_y, rotax_z).shape angle_rad = np.radians(angle_deg) rot_axes = np.array(np.broadcast_arrays( np.atleast_1d(rotax_x), np.atleast_1d(rotax_y), np.atleast_1d(rotax_z), )).astype(np.float64) d = rot_axes d /= np.linalg.norm(d, axis=0) W = np.zeros(rot_axes.shape[1:] + (3, 3)) W[..., 0, 1] = d[2] W[..., 0, 2] = -d[1] W[..., 1, 0] = -d[2] W[..., 1, 2] = d[0] W[..., 2, 0] = d[1] W[..., 2, 1] = -d[0] W2 = np.matmul(W, W) _a = angle_rad[..., np.newaxis, np.newaxis] eye = np.eye(3, dtype=np.float64) mtxs = eye + np.sin(_a) * W + 2. * np.sin(_a / 2.) ** 2 * W2 return mtxs.reshape(sh + (3, 3))
[docs]@utils.ranged_quantity_input( rotax_x=(None, None, apu.m), rotax_y=(None, None, apu.m), rotax_z=(None, None, apu.m), angle=(None, None, apu.deg), strip_input_units=True ) def rotmat_from_rotaxis(rotax_x, rotax_y, rotax_z, angle): ''' Construct rotation matrix from cartesian rotation axis and angle. Parameters ---------- rotax_{x,y,z} : `~astropy.units.Quantity` Cartesian components (x, y, z) of rotation axis [m] angle : `~astropy.units.Quantity` Rotation angle [deg] Returns ------- R : `~numpy.array` Rotation matrix [no units!] Notes ----- Broadcasting is supported. The rotation axis vector will be normalized in the function (otherwise it's magnitude would change the effective rotation angle). ''' return _rotmat_from_rotaxis(rotax_x, rotax_y, rotax_z, angle)
def _rotaxis_from_rotmat(R): # see Wikipedia (https://en.wikipedia.org/wiki/Axis–angle_representation) R = np.asarray(R) rot_angle = np.arccos( 0.5 * (np.trace(R, axis1=-2, axis2=-1) - 1) ) norm = -1. / 2. / np.sin(rot_angle) # different sign (vs. Wikipedia) rotax_x = norm * (R[..., 2, 1] - R[..., 1, 2]) rotax_y = norm * (R[..., 0, 2] - R[..., 2, 0]) rotax_z = norm * (R[..., 1, 0] - R[..., 0, 1]) return rotax_x, rotax_y, rotax_z, np.degrees(rot_angle)
[docs]@utils.ranged_quantity_input( output_unit=(apu.m, apu.m, apu.m, apu.deg) ) def rotaxis_from_rotmat(R): ''' Cartesian rotation axis and angle from rotation matrix. Parameters ---------- R : `~numpy.array` Rotation matrix [no units!] Returns ------- rotax_{x,y,z} : `~astropy.units.Quantity` Cartesian components (x, y, z) of rotation axis [m] angle : `~astropy.units.Quantity` Rotation angle [deg] Notes ----- Broadcasting is supported. The returned rotation axis vector is normalized (otherwise it's magnitude would change the effective rotation angle). ''' return _rotaxis_from_rotmat(R)
def _eulerangle_from_rotmat(R, etype='xyz'): ''' Euler angles from a rotation matrix. ''' # see https://www.geometrictools.com/Documentation/EulerAngles.pdf assert etype in ['xyz', 'zxz'] R = np.asarray(R) assert R.shape[-2:] == (3, 3) do_squeeze = False if R.ndim < 3: R = R.reshape((1, 3, 3)) do_squeeze = True alpha_1 = np.empty(R.shape[:-2], dtype=np.float64) alpha_2 = np.empty(R.shape[:-2], dtype=np.float64) alpha_3 = np.empty(R.shape[:-2], dtype=np.float64) if etype == 'xyz': mask1 = R[..., 2, 0] == 1 sub_R = R[mask1] alpha_1[mask1] = 0 alpha_2[mask1] = np.arctan2(-sub_R[:, 1, 2], sub_R[:, 1, 1]) alpha_3[mask1] = -np.pi / 2 mask2 = R[..., 2, 0] == -1 sub_R = R[mask2] alpha_1[mask2] = 0 alpha_2[mask2] = np.arctan2(-sub_R[:, 1, 2], sub_R[:, 1, 1]) alpha_3[mask2] = np.pi / 2 mask = (~mask1) & (~mask2) sub_R = R[mask] alpha_1[mask] = np.arctan2(sub_R[:, 2, 1], sub_R[:, 2, 2]) alpha_2[mask] = np.arcsin(-sub_R[:, 2, 0]) alpha_3[mask] = np.arctan2(sub_R[:, 1, 0], sub_R[:, 0, 0]) # why are the signs opposite to the EulerAngles.pdf document # (only for zyx, not for zxz?) alpha_1, alpha_2, alpha_3 = -alpha_1, -alpha_2, -alpha_3 elif etype == 'zxz': mask1 = R[..., 2, 2] == 1 sub_R = R[mask1] alpha_1[mask1] = 0 alpha_2[mask1] = 0 alpha_3[mask1] = np.arctan2(-sub_R[:, 0, 1], sub_R[:, 0, 0]) mask2 = R[..., 2, 2] == -1 sub_R = R[mask2] alpha_1[mask2] = 0 alpha_2[mask2] = np.pi alpha_3[mask2] = -np.arctan2(-sub_R[:, 1, 2], sub_R[:, 1, 1]) mask = (~mask1) & (~mask2) sub_R = R[mask] alpha_1[mask] = np.arctan2(sub_R[:, 2, 0], sub_R[:, 2, 1]) alpha_2[mask] = np.arccos(sub_R[:, 2, 2]) alpha_3[mask] = np.arctan2(sub_R[:, 0, 2], -sub_R[:, 1, 2]) # again, the above seems to be wrong; this fixes it: alpha_1 = np.pi - alpha_1 alpha_1 = (alpha_1 + np.pi) % (2 * np.pi) - np.pi alpha_3 = -np.pi - alpha_3 alpha_3 = (alpha_3 + np.pi) % (2 * np.pi) - np.pi alpha_1, alpha_2, alpha_3 = ( np.degrees(alpha_1), np.degrees(alpha_2), np.degrees(alpha_3) ) if do_squeeze: alpha_1, alpha_2, alpha_3 = ( alpha_1.squeeze(), alpha_2.squeeze(), alpha_3.squeeze() ) return alpha_3, alpha_2, alpha_1
[docs]@utils.ranged_quantity_input( output_unit=(apu.deg, apu.deg, apu.deg) ) def eulerangle_from_rotmat(R, etype='xyz'): ''' Cartesian rotation axis and angle from rotation matrix. Parameters ---------- R : `~numpy.array` Rotation matrix [no units!] etype : str, optional, 'xyz' or 'zxz' Desired Euler-angle ordering. (default: 'xyz') The ordering 'xyz' refers to the case where first a rotation about the x-axis is performed, then about y and last about z-axis, i.e., R = Rz(alpha3).Ry(alpha2).Rx(alpha1). Likewise for ordering 'zxz' it is R = Rz(alpha3).Rx(alpha2).Rz(alpha1) Returns ------- angle_{1, 2, 3} : `~astropy.units.Quantity` Euler rotation angles [deg] Notes ----- Broadcasting is supported. Note the returned Euler angles are not unique. ''' return _eulerangle_from_rotmat(R, etype=etype)