1# 2# This file is part of Healpy. 3# 4# Healpy is free software; you can redistribute it and/or modify 5# it under the terms of the GNU General Public License as published by 6# the Free Software Foundation; either version 2 of the License, or 7# (at your option) any later version. 8# 9# Healpy is distributed in the hope that it will be useful, 10# but WITHOUT ANY WARRANTY; without even the implied warranty of 11# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12# GNU General Public License for more details. 13# 14# You should have received a copy of the GNU General Public License 15# along with Healpy; if not, write to the Free Software 16# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA 17# 18# For more information about Healpy, see http://code.google.com/p/healpy 19# 20import numpy as np 21import six 22import warnings 23from . import pixelfunc 24from . import sphtfunc 25from ._sphtools import rotate_alm 26 27coordname = {"G": "Galactic", "E": "Ecliptic", "C": "Equatorial"} 28 29 30class ConsistencyWarning(Warning): 31 """Warns for a problem in the consistency of data 32 """ 33 34 pass 35 36 37if __name__ != "__main__": 38 warnings.filterwarnings("always", category=ConsistencyWarning, module=__name__) 39 40 41class Rotator(object): 42 """Rotation operator, including astronomical coordinate systems. 43 44 This class provides tools for spherical rotations. It is meant to be used 45 in the healpy library for plotting, and for this reason reflects the 46 convention used in the Healpix IDL library. 47 48 Parameters 49 ---------- 50 rot : None or sequence 51 Describe the rotation by its euler angle. See :func:`euler_matrix_new`. 52 coord : None or sequence of str 53 Describe the coordinate system transform. If *rot* is also given, the 54 coordinate transform is applied first, and then the rotation. 55 inv : bool 56 If True, the inverse rotation is defined. (Default: False) 57 deg : bool 58 If True, angles are assumed to be in degree. (Default: True) 59 eulertype : str 60 The Euler angle convention used. See :func:`euler_matrix_new`. 61 62 Attributes 63 ---------- 64 mat 65 coordin 66 coordout 67 coordinstr 68 coordoutstr 69 rots 70 coords 71 72 Examples 73 -------- 74 >>> r = Rotator(coord=['G','E']) # Transforms galactic to ecliptic coordinates 75 >>> theta_gal, phi_gal = np.pi/2., 0. 76 >>> theta_ecl, phi_ecl = r(theta_gal, phi_gal) # Apply the conversion 77 >>> print(theta_ecl) 78 1.66742286715 79 >>> print(phi_ecl) 80 -1.62596400306 81 >>> theta_ecl, phi_ecl = Rotator(coord='ge')(theta_gal, phi_gal) # In one line 82 >>> print(theta_ecl) 83 1.66742286715 84 >>> print(phi_ecl) 85 -1.62596400306 86 >>> vec_gal = np.array([1, 0, 0]) #Using vectors 87 >>> vec_ecl = r(vec_gal) 88 >>> print(vec_ecl) 89 [-0.05488249 -0.99382103 -0.09647625] 90 """ 91 92 ErrMessWrongPar = ( 93 "rot and coord must be single elements or " "sequence of same size." 94 ) 95 96 def __init__(self, rot=None, coord=None, inv=None, deg=True, eulertype="ZYX"): 97 """Create a rotator with given parameters. 98 - rot: a float, a tuple of 1,2 or 3 floats or a sequence of tuples. 99 If it is a sequence of tuple, it must have the same length as coord. 100 - coord: a string or a tuple of 1 or 2 strings or a sequence of tuple 101 If it is a sequence of tuple, it must have same length as rot. 102 - inv: whether to use inverse rotation or not 103 - deg: if True, angles in rot are assumed in degree (default: True) 104 - eulertype: the convention for Euler angles in rot. 105 Note: the coord system conversion is applied first, then the rotation. 106 """ 107 rot_is_seq = hasattr(rot, "__len__") and hasattr(rot[0], "__len__") 108 coord_is_seq = ( 109 hasattr(coord, "__len__") 110 and hasattr(coord[0], "__len__") 111 and type(coord[0]) is not str 112 ) 113 if rot_is_seq and coord_is_seq: 114 if len(rot) != len(coord): 115 raise ValueError(Rotator.ErrMessWrongPar) 116 else: 117 rots = rot 118 coords = coord 119 elif (rot_is_seq or coord_is_seq) and (rot is not None and coord is not None): 120 raise ValueError(Rotator.ErrMessWrongPar) 121 else: 122 rots = [rot] 123 coords = [coord] 124 inv_is_seq = hasattr(inv, "__len__") 125 if inv_is_seq: 126 if len(inv) != len(rots): 127 raise ValueError("inv must have same length as rot and/or coord") 128 invs = inv 129 else: 130 invs = [inv] * len(rots) 131 # check the argument and normalize them 132 if eulertype in ["ZYX", "X", "Y"]: 133 self._eultype = eulertype 134 else: 135 self._eultype = "ZYX" 136 self._rots = [] 137 self._coords = [] 138 self._invs = [] 139 for r, c, i in zip(rots, coords, invs): 140 rn = normalise_rot(r, deg=deg) 141 # if self._eultype in ['X','Y']: 142 # rn[1] = -rn[1] 143 cn = normalise_coord(c) 144 self._rots.append(rn) # append(rn) or insert(0, rn) ? 145 self._coords.append(cn) # append(cn) or insert(0, cn) ? 146 self._invs.append(bool(i)) 147 if not self.consistent: 148 warnings.warn( 149 "The chain of coord system rotations is not consistent", 150 category=ConsistencyWarning, 151 ) 152 self._update_matrix() 153 154 def _update_matrix(self): 155 self._matrix = np.identity(3) 156 self._do_rotation = False 157 for r, c, i in zip(self._rots, self._coords, self._invs): 158 rotmat, do_rot, rotnorm = get_rotation_matrix(r, eulertype=self._eultype) 159 convmat, do_conv, coordnorm = get_coordconv_matrix(c) 160 r = np.dot(rotmat, convmat) 161 if i: 162 r = r.T 163 self._matrix = np.dot(self._matrix, r) 164 self._do_rotation = self._do_rotation or (do_rot or do_conv) 165 166 def _is_coords_consistent(self): 167 for c, i in zip(self._coords, self._invs): 168 break 169 for cnext, inext in zip(self._coords[1:], self._invs[1:]): 170 if c[i] != cnext[not inext]: 171 return False 172 c, i = cnext, inext 173 return True 174 175 consistent = property( 176 _is_coords_consistent, doc="consistency of the coords transform chain" 177 ) 178 179 def __eq__(self, a): 180 if type(a) is not type(self): 181 return False 182 # compare the _rots 183 v = [np.allclose(x, y, rtol=0, atol=1e-15) for x, y in zip(self._rots, a._rots)] 184 return ( 185 np.array(v).all() 186 and (self._coords == a._coords) 187 and (self._invs == a._invs) 188 ) 189 190 def __call__(self, *args, **kwds): 191 """Use the rotator to rotate either spherical coordinates (theta, phi) 192 or a vector (x,y,z). You can use lonla keyword to use longitude, latitude 193 (in degree) instead of theta, phi (in radian). In this case, returns 194 longitude, latitude in degree. 195 196 Accepted forms: 197 198 r(x,y,z) # x,y,z either scalars or arrays 199 r(theta,phi) # theta, phi scalars or arrays 200 r(lon,lat,lonlat=True) # lon, lat scalars or arrays 201 r(vec) # vec 1-D array with 3 elements, or 2-D array 3xN 202 r(direction) # direction 1-D array with 2 elements, or 2xN array 203 204 Parameters 205 ---------- 206 vec_or_dir : array or multiple arrays 207 The direction to rotate. See above for accepted formats. 208 lonlat : bool, optional 209 If True, assumes the input direction is longitude/latitude in degrees. 210 Otherwise, assumes co-latitude/longitude in radians. Default: False 211 inv : bool, optional 212 If True, applies the inverse rotation. Default: False. 213 """ 214 if kwds.pop("inv", False): 215 m = self._matrix.T 216 else: 217 m = self._matrix 218 lonlat = kwds.pop("lonlat", False) 219 if len(args) == 1: 220 arg = args[0] 221 if not hasattr(arg, "__len__") or len(arg) < 2 or len(arg) > 3: 222 raise TypeError("Argument must be a sequence of 2 or 3 " "elements") 223 if len(arg) == 2: 224 return rotateDirection( 225 m, arg[0], arg[1], self._do_rotation, lonlat=lonlat 226 ) 227 else: 228 return rotateVector(m, arg[0], arg[1], arg[2], self._do_rotation) 229 elif len(args) == 2: 230 return rotateDirection( 231 m, args[0], args[1], self._do_rotation, lonlat=lonlat 232 ) 233 elif len(args) == 3: 234 return rotateVector(m, args[0], args[1], args[2], self._do_rotation) 235 else: 236 raise TypeError("Either 1, 2 or 3 arguments accepted") 237 238 def __mul__(self, a): 239 """Composition of rotation. 240 """ 241 if not isinstance(a, Rotator): 242 raise TypeError( 243 "A Rotator can only multiply another Rotator " 244 "(composition of rotations)" 245 ) 246 rots = self._rots + a._rots 247 coords = self._coords + a._coords 248 invs = self._invs + a._invs 249 return Rotator(rot=rots, coord=coords, inv=invs, deg=False) 250 251 def __rmul__(self, b): 252 if not isinstance(b, Rotator): 253 raise TypeError( 254 "A Rotator can only be multiplied by another Rotator " 255 "(composition of rotations)" 256 ) 257 rots = b._rots + self._rots 258 coords = b._coords + self._coords 259 invs = self._invs + a._invs 260 return Rotator(rot=rots, coord=coords, inv=invs, deg=False) 261 262 def __nonzero__(self): 263 return self._do_rotation 264 265 def get_inverse(self): 266 rots = self._rots[::-1] 267 coords = self._coords[::-1] 268 invs = [not i for i in self._invs[::-1]] 269 return Rotator(rot=rots, coord=coords, inv=invs, deg=False) 270 271 # I = property(get_inverse,doc='Return a new rotator representing the ' 272 # 'inverse rotation') 273 274 def I(self, *args, **kwds): 275 """Rotate the given vector or direction using the inverse matrix. 276 rot.I(vec) <==> rot(vec,inv=True) 277 """ 278 kwds["inv"] = True 279 return self.__call__(*args, **kwds) 280 281 @property 282 def mat(self): 283 """The matrix representing the rotation. 284 """ 285 return np.matrix(self._matrix) 286 287 @property 288 def coordin(self): 289 """The input coordinate system. 290 """ 291 if not self.consistent: 292 return None 293 for c, i in zip(self._coords, self._invs): 294 pass 295 return c[i] 296 297 @property 298 def coordout(self): 299 """The output coordinate system. 300 """ 301 if not self.consistent: 302 return None 303 for c, i in zip(self._coords, self._invs): 304 pass 305 return c[not i] 306 307 @property 308 def coordinstr(self): 309 """The input coordinate system in str. 310 """ 311 return coordname.get(self.coordin, "") 312 313 @property 314 def coordoutstr(self): 315 """The output coordinate system in str. 316 """ 317 return coordname.get(self.coordout, "") 318 319 @property 320 def rots(self): 321 """The sequence of rots defining the rotation. 322 """ 323 return self._rots 324 325 @property 326 def coords(self): 327 """The sequence of coords defining the rotation. 328 """ 329 return self._coords 330 331 def do_rot(self, i): 332 """Returns True if rotation is not (close to) identity. 333 """ 334 return not np.allclose(self.rots[i], np.zeros(3), rtol=0., atol=1.e-15) 335 336 def angle_ref(self, *args, **kwds): 337 """Compute the angle between transverse reference direction of initial and final frames 338 339 For example, if angle of polarisation is psi in initial frame, it will be psi+angle_ref in final 340 frame. 341 342 Parameters 343 ---------- 344 dir_or_vec : array 345 Direction or vector (see Rotator.__call__) 346 lonlat: bool, optional 347 If True, assume input is longitude,latitude in degrees. Otherwise, 348 theta,phi in radian. Default: False 349 inv : bool, optional 350 If True, use the inverse transforms. Default: False 351 352 Returns 353 ------- 354 angle : float, scalar or array 355 Angle in radian (a scalar or an array if input is a sequence of direction/vector) 356 """ 357 R = self 358 lonlat = kwds.get("lonlat", False) 359 inv = kwds.get("inv", False) 360 if len(args) == 1: 361 arg = args[0] 362 if not hasattr(arg, "__len__") or len(arg) < 2 or len(arg) > 3: 363 raise TypeError("Argument must be a sequence of 2 or 3 " "elements") 364 if len(arg) == 2: 365 v = dir2vec(arg[0], arg[1], lonlat=lonlat) 366 else: 367 v = arg 368 elif len(args) == 2: 369 v = dir2vec(args[0], args[1], lonlat=lonlat) 370 elif len(args) == 3: 371 v = args 372 else: 373 raise TypeError("Either 1, 2 or 3 arguments accepted") 374 vp = R(v, inv=inv) 375 north_pole = R([0., 0., 1.], inv=inv) 376 sinalpha = north_pole[0] * vp[1] - north_pole[1] * vp[0] 377 cosalpha = north_pole[2] - vp[2] * np.dot(north_pole, vp) 378 return np.arctan2(sinalpha, cosalpha) 379 380 def rotate_alm(self, alm, lmax=None, mmax=None): 381 """Rotate Alms with the transform defined in the Rotator object 382 383 see the docstring of the rotate_alm function defined 384 in the healpy package, this function **returns** the rotated alms, 385 does not rotate in place""" 386 387 rotated_alm = alm.copy() # rotate_alm works inplace 388 rotate_alm(rotated_alm, matrix=self.mat, lmax=lmax, mmax=mmax) 389 return rotated_alm 390 391 def rotate_map_alms(self, m, use_pixel_weights=True, lmax=None, mmax=None): 392 """Rotate a HEALPix map to a new reference frame in spherical harmonics space 393 394 This is generally the best strategy to rotate/change reference frame of maps. 395 If the input map is band-limited, i.e. it can be represented exactly by 396 a spherical harmonics transform under a specific lmax, the map rotation 397 will be invertible. 398 399 400 Parameters 401 ---------- 402 m : np.ndarray 403 Input map, 1 map is considered I, 2 maps:[Q,U], 3 maps:[I,Q,U] 404 use_pixel_weights : bool, optional 405 Use pixel weights in map2alm 406 407 Returns 408 ------- 409 m_rotated : np.ndarray 410 Map in the new reference frame 411 """ 412 alm = sphtfunc.map2alm( 413 m, use_pixel_weights=use_pixel_weights, lmax=lmax, mmax=mmax 414 ) 415 rotated_alm = self.rotate_alm(alm, lmax=lmax, mmax=mmax) 416 return sphtfunc.alm2map( 417 rotated_alm, lmax=lmax, mmax=mmax, nside=pixelfunc.get_nside(m) 418 ) 419 420 def rotate_map_pixel(self, m): 421 """Rotate a HEALPix map to a new reference frame in pixel space 422 423 It is generally better to rotate in spherical harmonics space, see 424 the rotate_map_alms method. A case where pixel space rotation is 425 better is for heavily masked maps where the spherical harmonics 426 transform is not well defined. 427 This function first rotates the pixels centers of the new reference 428 frame to the original reference frame, then uses hp.get_interp_val 429 to interpolate bilinearly the pixel values, finally fixes Q and U 430 polarization by the modification to the psi angle caused by 431 the Rotator using Rotator.angle_ref. 432 Due to interpolation, this function generally suppresses the signal at 433 high angular scales. 434 435 Parameters 436 ---------- 437 m : np.ndarray 438 Input map, 1 map is considered I, 2 maps:[Q,U], 3 maps:[I,Q,U] 439 440 Returns 441 ------- 442 m_rotated : np.ndarray 443 Map in the new reference frame 444 445 """ 446 if pixelfunc.maptype(m) == 0: # a single map is converted to a list 447 m = [m] 448 npix = len(m[0]) 449 nside = pixelfunc.npix2nside(npix) 450 theta_pix_center, phi_pix_center = pixelfunc.pix2ang( 451 nside=nside, ipix=np.arange(npix) 452 ) 453 454 # Rotate the pixels center of the new reference frame to the original frame 455 theta_pix_center_rot, phi_pix_center_rot = self.I( 456 theta_pix_center, phi_pix_center 457 ) 458 459 # Interpolate the original map to the pixels centers in the new ref frame 460 m_rotated = [ 461 pixelfunc.get_interp_val(each, theta_pix_center_rot, phi_pix_center_rot) 462 for each in m 463 ] 464 465 # Rotate polarization 466 if len(m_rotated) > 1: 467 # Create a complex map from QU and apply the rotation in psi due to the rotation 468 # Slice from the end of the array so that it works both for QU and IQU 469 L_map = (m_rotated[-2] + m_rotated[-1] * 1j) * np.exp( 470 1j * 2 * self.angle_ref(theta_pix_center_rot, phi_pix_center_rot) 471 ) 472 473 # Overwrite the Q and U maps with the correct values 474 m_rotated[-2] = np.real(L_map) 475 m_rotated[-1] = np.imag(L_map) 476 else: 477 m_rotated = m_rotated[0] 478 479 return m_rotated 480 481 def __repr__(self): 482 return ( 483 "[ " 484 + ", ".join([str(self._coords), str(self._rots), str(self._invs)]) 485 + " ]" 486 ) 487 488 __str__ = __repr__ 489 490 491################################################################ 492# 493# Helpers function for rotation 494# used in the Rotator class. 495 496 497def rotateVector(rotmat, vec, vy=None, vz=None, do_rot=True): 498 """Rotate a vector (or a list of vectors) using the rotation matrix 499 given as first argument. 500 501 Parameters 502 ---------- 503 rotmat : float, array-like shape (3,3) 504 The rotation matrix 505 vec : float, scalar or array-like 506 The vector to transform (shape (3,) or (3,N)), 507 or x component (scalar or shape (N,)) if vy and vz are given 508 vy : float, scalar or array-like, optional 509 The y component of the vector (scalar or shape (N,)) 510 vz : float, scalar or array-like, optional 511 The z component of the vector (scalar or shape (N,)) 512 do_rot : bool, optional 513 if True, really perform the operation, if False do nothing. 514 515 Returns 516 ------- 517 vec : float, array 518 The component of the rotated vector(s). 519 520 See Also 521 -------- 522 Rotator 523 """ 524 if vy is None and vz is None: 525 if do_rot: 526 return np.tensordot(rotmat, vec, axes=(1, 0)) 527 else: 528 return vec 529 elif vy is not None and vz is not None: 530 if do_rot: 531 return np.tensordot(rotmat, np.array([vec, vy, vz]), axes=(1, 0)) 532 else: 533 return vec, vy, vz 534 else: 535 raise TypeError("You must give either vec only or vec, vy " "and vz parameters") 536 537 538def rotateDirection(rotmat, theta, phi=None, do_rot=True, lonlat=False): 539 """Rotate the vector described by angles theta,phi using the rotation matrix 540 given as first argument. 541 542 Parameters 543 ---------- 544 rotmat : float, array-like shape (3,3) 545 The rotation matrix 546 theta : float, scalar or array-like 547 The angle theta (scalar or shape (N,)) 548 or both angles (scalar or shape (2, N)) if phi is not given. 549 phi : float, scalar or array-like, optionnal 550 The angle phi (scalar or shape (N,)). 551 do_rot : bool, optional 552 if True, really perform the operation, if False do nothing. 553 lonlat : bool 554 If True, input angles are assumed to be longitude and latitude in degree, 555 otherwise, they are co-latitude and longitude in radians. 556 557 Returns 558 ------- 559 angles : float, array 560 The angles of describing the rotated vector(s). 561 562 See Also 563 -------- 564 Rotator 565 """ 566 vx, vy, vz = rotateVector(rotmat, dir2vec(theta, phi, lonlat=lonlat), do_rot=do_rot) 567 return vec2dir(vx, vy, vz, lonlat=lonlat) 568 569 570def vec2dir(vec, vy=None, vz=None, lonlat=False): 571 """Transform a vector to angle given by theta,phi. 572 573 Parameters 574 ---------- 575 vec : float, scalar or array-like 576 The vector to transform (shape (3,) or (3,N)), 577 or x component (scalar or shape (N,)) if vy and vz are given 578 vy : float, scalar or array-like, optional 579 The y component of the vector (scalar or shape (N,)) 580 vz : float, scalar or array-like, optional 581 The z component of the vector (scalar or shape (N,)) 582 lonlat : bool, optional 583 If True, return angles will be longitude and latitude in degree, 584 otherwise, angles will be longitude and co-latitude in radians (default) 585 586 Returns 587 ------- 588 angles : float, array 589 The angles (unit depending on *lonlat*) in an array of 590 shape (2,) (if scalar input) or (2, N) 591 592 See Also 593 -------- 594 :func:`dir2vec`, :func:`pixelfunc.ang2vec`, :func:`pixelfunc.vec2ang` 595 """ 596 if np.any(np.isnan(vec)): 597 return np.nan, np.nan 598 if vy is None and vz is None: 599 vx, vy, vz = vec 600 elif vy is not None and vz is not None: 601 vx = vec 602 else: 603 raise TypeError("You must either give both vy and vz or none of them") 604 r = np.sqrt(vx ** 2 + vy ** 2 + vz ** 2) 605 ang = np.empty((2, r.size)) 606 ang[0, :] = np.arccos(vz / r) 607 ang[1, :] = np.arctan2(vy, vx) 608 if lonlat: 609 ang = np.degrees(ang) 610 np.negative(ang[0, :], ang[0, :]) 611 ang[0, :] += 90. 612 return ang[::-1, :].squeeze() 613 else: 614 return ang.squeeze() 615 616 617def dir2vec(theta, phi=None, lonlat=False): 618 """Transform a direction theta,phi to a unit vector. 619 620 Parameters 621 ---------- 622 theta : float, scalar or array-like 623 The angle theta (scalar or shape (N,)) 624 or both angles (scalar or shape (2, N)) if phi is not given. 625 phi : float, scalar or array-like, optionnal 626 The angle phi (scalar or shape (N,)). 627 lonlat : bool 628 If True, input angles are assumed to be longitude and latitude in degree, 629 otherwise, they are co-latitude and longitude in radians. 630 631 Returns 632 ------- 633 vec : array 634 The vector(s) corresponding to given angles, shape is (3,) or (3, N). 635 636 See Also 637 -------- 638 :func:`vec2dir`, :func:`pixelfunc.ang2vec`, :func:`pixelfunc.vec2ang` 639 """ 640 if phi is None: 641 theta, phi = theta 642 if lonlat: 643 lon, lat = theta, phi 644 theta, phi = np.pi / 2. - np.radians(lat), np.radians(lon) 645 ct, st, cp, sp = np.cos(theta), np.sin(theta), np.cos(phi), np.sin(phi) 646 vec = np.empty((3, ct.size), np.float64) 647 vec[0, :] = st * cp 648 vec[1, :] = st * sp 649 vec[2, :] = ct 650 return vec.squeeze() 651 652 653def angdist(dir1, dir2, lonlat=False): 654 """Returns the angular distance between dir1 and dir2. 655 656 Parameters 657 ---------- 658 dir1, dir2 : float, array-like 659 The directions between which computing the angular distance. 660 Angular if len(dir) == 2 or vector if len(dir) == 3. 661 See *lonlat* for unit 662 lonlat : bool, scalar or sequence 663 If True, angles are assumed to be longitude and latitude in degree, 664 otherwise they are interpreted as colatitude and longitude in radian. 665 If a sequence, lonlat[0] applies to dir1 and lonlat[1] applies to dir2. 666 667 Returns 668 ------- 669 angles : float, scalar or array-like 670 The angle(s) between dir1 and dir2 in radian. 671 672 Examples 673 -------- 674 >>> import healpy as hp 675 >>> hp.rotator.angdist([.2,0], [.2, 1e-6]) 676 array([ 1.98669331e-07]) 677 """ 678 if hasattr(lonlat, "__len__") and len(lonlat) == 2: 679 lonlat1, lonlat2 = lonlat 680 else: 681 lonlat1 = lonlat2 = lonlat 682 dir1 = np.asarray(dir1) 683 dir2 = np.asarray(dir2) 684 if dir1.ndim == 2: 685 if dir1.shape[0] == 2: # theta, phi -> vec 686 vec1 = dir2vec(dir1, lonlat=lonlat1) 687 else: 688 vec1 = np.reshape(dir1, (3, -1)) 689 vec1 = normalize_vec(vec1) 690 elif dir1.ndim == 1: 691 if dir1.shape[0] == 2: # theta, phi -> vec 692 vec1 = np.reshape(dir2vec(dir1, lonlat=lonlat1), (3, 1)) 693 else: 694 vec1 = np.reshape(dir1, (3, 1)) 695 vec1 = normalize_vec(vec1) 696 if dir2.ndim == 2: 697 if dir2.shape[0] == 2: # theta, phi -> vec 698 vec2 = dir2vec(dir2, lonlat=lonlat2) 699 else: 700 vec2 = np.reshape(dir2, (3, -1)) 701 vec2 = normalize_vec(vec2) 702 elif dir2.ndim == 1: 703 if dir2.shape[0] == 2: # theta, phi -> vec 704 vec2 = np.reshape(dir2vec(dir2, lonlat=lonlat2), (3, 1)) 705 else: 706 vec2 = np.reshape(dir2, (3, 1)) 707 vec2 = normalize_vec(vec2) 708 # compute vec product 709 vec_prod = np.sqrt((np.cross(vec1.T, vec2.T) ** 2).sum(axis=1)) 710 # compute scalar product 711 scal_prod = (vec1 * vec2).sum(axis=0) 712 713 return np.arctan2(vec_prod, scal_prod) 714 715 716def normalize_vec(vec): 717 """Normalize the vector(s) *vec* (in-place if it is a ndarray). 718 719 Parameters 720 ---------- 721 vec : float, array-like of shape (D,) or (D, N) 722 The D-vector(s) to normalize. 723 724 Returns 725 ------- 726 vec_normed : float, array 727 Normalized vec, shape (D,) or (D, N) 728 """ 729 vec = np.array(vec, np.float64) 730 r = np.sqrt(np.sum(vec ** 2, axis=0)) 731 vec /= r 732 return vec 733 734 735####################################################### 736# 737# Manage the coord system conventions 738# 739 740 741def check_coord(c): 742 """Check if parameter is a valid coord system. 743 Raise a TypeError exception if it is not, otherwise returns the normalized 744 coordinate system name. 745 """ 746 if c is None: 747 return c 748 if not isinstance(c, six.string_types): 749 raise TypeError( 750 "Coordinate must be a string (G[alactic]," 751 " E[cliptic], C[elestial]" 752 " or Equatorial=Celestial)" 753 ) 754 if c[0].upper() == "G": 755 x = "G" 756 elif c[0].upper() == "E" and c != "Equatorial": 757 x = "E" 758 elif c[0].upper() == "C" or c == "Equatorial": 759 x = "C" 760 else: 761 raise ValueError( 762 "Wrong coordinate (either G[alactic]," 763 " E[cliptic], C[elestial]" 764 " or Equatorial=Celestial)" 765 ) 766 return x 767 768 769def normalise_coord(coord): 770 """Normalise the coord argument. 771 Coord sys are either 'E','G', 'C' or 'X' if undefined. 772 773 Input: either a string or a sequence of string. 774 775 Output: a tuple of two strings, each being one of the norm coord sys name 776 above. 777 778 eg, 'E' -> ['E','E'], ['Ecliptic','G'] -> ['E','G'] 779 None -> ['X','X'] etc. 780 """ 781 coord_norm = [] 782 if coord is None: 783 coord = (None, None) 784 coord = tuple(coord) 785 if len(coord) > 2: 786 raise TypeError( 787 "Coordinate must be a string (G[alactic]," 788 " E[cliptic] or C[elestial])" 789 " or a sequence of 2 strings" 790 ) 791 for x in coord: 792 coord_norm.append(check_coord(x)) 793 if len(coord_norm) < 2: 794 coord_norm.append(coord_norm[0]) 795 return tuple(coord_norm) 796 797 798def normalise_rot(rot, deg=False): 799 """Return rot possibly completed with zeroes to reach size 3. 800 If rot is None, return a vector of 0. 801 If deg is True, convert from degree to radian, otherwise assume input 802 is in radian. 803 """ 804 if deg: 805 convert = np.pi / 180. 806 else: 807 convert = 1. 808 if rot is None: 809 rot = np.zeros(3) 810 else: 811 rot = np.array(rot, np.float64).flatten() * convert 812 rot.resize(3) 813 return rot 814 815 816def get_rotation_matrix(rot, deg=False, eulertype="ZYX"): 817 """Return the rotation matrix corresponding to angles given in rot. 818 819 Usage: matrot,do_rot,normrot = get_rotation_matrix(rot) 820 821 Input: 822 - rot: either None, an angle or a tuple of 1,2 or 3 angles 823 corresponding to Euler angles. 824 Output: 825 - matrot: 3x3 rotation matrix 826 - do_rot: True if rotation is not identity, False otherwise 827 - normrot: the normalized version of the input rot. 828 """ 829 rot = normalise_rot(rot, deg=deg) 830 if not np.allclose(rot, np.zeros(3), rtol=0., atol=1.e-15): 831 do_rot = True 832 else: 833 do_rot = False 834 if eulertype == "X": 835 matrot = euler_matrix_new(rot[0], -rot[1], rot[2], X=True) 836 elif eulertype == "Y": 837 matrot = euler_matrix_new(rot[0], -rot[1], rot[2], Y=True) 838 else: 839 matrot = euler_matrix_new(rot[0], -rot[1], rot[2], ZYX=True) 840 841 return matrot, do_rot, rot 842 843 844def get_coordconv_matrix(coord): 845 """Return the rotation matrix corresponding to coord systems given 846 in coord. 847 848 Usage: matconv,do_conv,normcoord = get_coordconv_matrix(coord) 849 850 Input: 851 - coord: a tuple with initial and final coord systems. 852 See normalise_coord. 853 Output: 854 - matconv: the euler matrix for coord sys conversion 855 - do_conv: True if matconv is not identity, False otherwise 856 - normcoord: the tuple of initial and final coord sys. 857 858 History: adapted from CGIS IDL library. 859 """ 860 861 coord_norm = normalise_coord(coord) 862 863 if coord_norm[0] == coord_norm[1]: 864 matconv = np.identity(3) 865 do_conv = False 866 else: 867 eps = 23.452294 - 0.0130125 - 1.63889E-6 + 5.02778E-7 868 eps = eps * np.pi / 180. 869 870 # ecliptic to galactic 871 e2g = np.array( 872 [ 873 [-0.054882486, -0.993821033, -0.096476249], 874 [0.494116468, -0.110993846, 0.862281440], 875 [-0.867661702, -0.000346354, 0.497154957], 876 ] 877 ) 878 879 # ecliptic to equatorial 880 e2q = np.array( 881 [ 882 [1., 0., 0.], 883 [0., np.cos(eps), -1. * np.sin(eps)], 884 [0., np.sin(eps), np.cos(eps)], 885 ] 886 ) 887 888 # galactic to ecliptic 889 g2e = np.linalg.inv(e2g) 890 891 # galactic to equatorial 892 g2q = np.dot(e2q, g2e) 893 894 # equatorial to ecliptic 895 q2e = np.linalg.inv(e2q) 896 897 # equatorial to galactic 898 q2g = np.dot(e2g, q2e) 899 900 if coord_norm == ("E", "G"): 901 matconv = e2g 902 elif coord_norm == ("G", "E"): 903 matconv = g2e 904 elif coord_norm == ("E", "C"): 905 matconv = e2q 906 elif coord_norm == ("C", "E"): 907 matconv = q2e 908 elif coord_norm == ("C", "G"): 909 matconv = q2g 910 elif coord_norm == ("G", "C"): 911 matconv = g2q 912 else: 913 raise ValueError("Wrong coord transform :", coord_norm) 914 do_conv = True 915 916 return matconv, do_conv, coord_norm 917 918 919################################################### 920## ## 921## euler functions ## 922## ## 923###### ####### 924 925 926def euler(ai, bi, select, FK4=0): 927 """ 928 NAME: 929 euler 930 PURPOSE: 931 Transform between Galactic, celestial, and ecliptic coordinates. 932 EXPLANATION: 933 Use the procedure ASTRO to use this routine interactively 934 935 CALLING SEQUENCE: 936 EULER, AI, BI, AO, BO, [ SELECT, /FK4, SELECT = ] 937 938 INPUTS: 939 AI - Input Longitude in DEGREES, scalar or vector. If only two 940 parameters are supplied, then AI and BI will be modified 941 to contain the output longitude and latitude. 942 BI - Input Latitude in DEGREES 943 944 OPTIONAL INPUT: 945 SELECT - Integer (1-6) specifying type of coordinate 946 transformation. 947 948 SELECT From To | SELECT From To 949 1 RA-Dec (2000) Galactic | 4 Ecliptic RA-Dec 950 2 Galactic RA-DEC | 5 Ecliptic Galactic 951 3 RA-Dec Ecliptic | 6 Galactic Ecliptic 952 953 If not supplied as a parameter or keyword, then EULER will prompt 954 for the value of SELECT 955 Celestial coordinates (RA, Dec) should be given in equinox J2000 956 unless the /FK4 keyword is set. 957 OUTPUTS: 958 AO - Output Longitude in DEGREES 959 BO - Output Latitude in DEGREES 960 961 INPUT KEYWORD: 962 /FK4 - If this keyword is set and non-zero, then input and output 963 celestial and ecliptic coordinates should be given in 964 equinox B1950. 965 /SELECT - The coordinate conversion integer (1-6) may 966 alternatively be specified as a keyword 967 NOTES: 968 EULER was changed in December 1998 to use J2000 coordinates as the 969 default, ** and may be incompatible with earlier versions***. 970 REVISION HISTORY: 971 Written W. Landsman, February 1987 972 Adapted from Fortran by Daryl Yentis NRL 973 Converted to IDL V5.0 W. Landsman September 1997 974 Made J2000 the default, added /FK4 keyword 975 W. Landsman December 1998 976 Add option to specify SELECT as a keyword W. Landsman March 2003 977 Converted to python by K. Ganga December 2007 978 """ 979 980 # npar = N_params() 981 # if npar LT 2 then begin 982 # print,'Syntax - EULER, AI, BI, A0, B0, [ SELECT, /FK4, SELECT= ]' 983 # print,' AI,BI - Input longitude,latitude in degrees' 984 # print,' AO,BO - Output longitude, latitude in degrees' 985 # print,' SELECT - Scalar (1-6) specifying transformation type' 986 # return 987 # endif 988 989 PI = np.pi 990 twopi = 2.0 * PI 991 fourpi = 4.0 * PI 992 deg_to_rad = 180.0 / PI 993 # 994 # ; J2000 coordinate conversions are based on the following constants 995 # ; (see the Hipparcos explanatory supplement). 996 # ; eps = 23.4392911111 # Obliquity of the ecliptic 997 # ; alphaG = 192.85948d Right Ascension of Galactic North Pole 998 # ; deltaG = 27.12825d Declination of Galactic North Pole 999 # ; lomega = 32.93192d Galactic longitude of celestial equator 1000 # ; alphaE = 180.02322d Ecliptic longitude of Galactic North Pole 1001 # ; deltaE = 29.811438523d Ecliptic latitude of Galactic North Pole 1002 # ; Eomega = 6.3839743d Galactic longitude of ecliptic equator 1003 # 1004 if FK4 == 1: 1005 1006 equinox = "(B1950)" 1007 psi = [ 1008 0.57595865315, 1009 4.9261918136, 1010 0.00000000000, 1011 0.0000000000, 1012 0.11129056012, 1013 4.7005372834, 1014 ] 1015 stheta = [ 1016 0.88781538514, 1017 -0.88781538514, 1018 0.39788119938, 1019 -0.39788119938, 1020 0.86766174755, 1021 -0.86766174755, 1022 ] 1023 ctheta = [ 1024 0.46019978478, 1025 0.46019978478, 1026 0.91743694670, 1027 0.91743694670, 1028 0.49715499774, 1029 0.49715499774, 1030 ] 1031 phi = [ 1032 4.9261918136, 1033 0.57595865315, 1034 0.0000000000, 1035 0.00000000000, 1036 4.7005372834, 1037 0.11129056012, 1038 ] 1039 else: 1040 1041 equinox = "(J2000)" 1042 psi = [ 1043 0.57477043300, 1044 4.9368292465, 1045 0.00000000000, 1046 0.0000000000, 1047 0.11142137093, 1048 4.71279419371, 1049 ] 1050 stheta = [ 1051 0.88998808748, 1052 -0.88998808748, 1053 0.39777715593, 1054 -0.39777715593, 1055 0.86766622025, 1056 -0.86766622025, 1057 ] 1058 ctheta = [ 1059 0.45598377618, 1060 0.45598377618, 1061 0.91748206207, 1062 0.91748206207, 1063 0.49714719172, 1064 0.49714719172, 1065 ] 1066 phi = [ 1067 4.9368292465, 1068 0.57477043300, 1069 0.0000000000, 1070 0.00000000000, 1071 4.71279419371, 1072 0.11142137093, 1073 ] 1074 # 1075 i = select - 1 # IDL offset 1076 a = ai / deg_to_rad - phi[i] 1077 b = bi / deg_to_rad 1078 sb = np.sin(b) 1079 cb = np.cos(b) 1080 cbsa = cb * np.sin(a) 1081 b = -stheta[i] * cbsa + ctheta[i] * sb 1082 # bo = math.asin(where(b<1.0, b, 1.0)*deg_to_rad) 1083 bo = np.arcsin(b) * deg_to_rad 1084 # 1085 a = np.arctan2(ctheta[i] * cbsa + stheta[i] * sb, cb * np.cos(a)) 1086 ao = np.fmod((a + psi[i] + fourpi), twopi) * deg_to_rad 1087 return ao, bo 1088 1089 1090def euler_matrix_new(a1, a2, a3, X=True, Y=False, ZYX=False, deg=False): 1091 """ 1092 NAME: 1093 euler_matrix_new 1094 1095 PURPOSE: 1096 computes the Euler matrix of an arbitrary rotation described 1097 by 3 Euler angles 1098 correct bugs present in Euler_Matrix 1099 1100 CALLING SEQUENCE: 1101 result = euler_matrix_new (a1, a2, a3 [,X, Y, ZYX, DEG ]) 1102 1103 INPUTS: 1104 a1, a2, a3 = Euler angles, scalar 1105 (in radian by default, in degree if DEG is set) 1106 all the angles are measured counterclockwise 1107 correspond to x, y, zyx-conventions (see Goldstein) 1108 the default is x 1109 1110 KEYWORD PARAMETERS: 1111 DEG : if set the angle are measured in degree 1112 1113 X : rotation a1 around original Z 1114 rotation a2 around interm X 1115 rotation a3 around final Z 1116 DEFAULT, classical mechanics convention 1117 1118 Y : rotation a1 around original Z 1119 rotation a2 around interm Y 1120 rotation a3 around final Z 1121 quantum mechanics convention (override X) 1122 1123 ZYX : rotation a1 around original Z 1124 rotation a2 around interm Y 1125 rotation a3 around final X 1126 aeronautics convention (override X) 1127 1128 * these last three keywords are obviously mutually exclusive * 1129 1130 OUTPUTS: 1131 result is a 3x3 matrix 1132 1133 USAGE: 1134 if vec is an Nx3 array containing N 3D vectors, 1135 vec # euler_matrix_new(a1,a2,a3,/Y) will be the rotated vectors 1136 1137 1138 MODIFICATION HISTORY: 1139 March 2002, EH, Caltech, rewritting of euler_matrix 1140 1141 convention euler_matrix_new euler_matrix 1142 X: M_new(a,b,c,/X) = M_old(-a,-b,-c,/X) = Transpose( M_old(c, b, a,/X)) 1143 Y: M_new(a,b,c,/Y) = M_old(-a, b,-c,/Y) = Transpose( M_old(c,-b, a,/Y)) 1144 ZYX: M_new(a,b,c,/Z) = M_old(-a, b,-c,/Z) 1145 """ 1146 1147 t_k = 0 1148 if ZYX: 1149 t_k = t_k + 1 1150 # if X: t_k = t_k + 1 1151 if Y: 1152 t_k = t_k + 1 1153 if t_k > 1: 1154 raise ValueError("Choose either X, Y or ZYX convention") 1155 1156 convert = 1.0 1157 if deg: 1158 convert = np.pi / 180. 1159 1160 c1 = np.cos(a1 * convert) 1161 s1 = np.sin(a1 * convert) 1162 c2 = np.cos(a2 * convert) 1163 s2 = np.sin(a2 * convert) 1164 c3 = np.cos(a3 * convert) 1165 s3 = np.sin(a3 * convert) 1166 1167 if ZYX: 1168 m1 = np.array([[c1, -s1, 0], [s1, c1, 0], [0, 0, 1]]) # around z 1169 1170 m2 = np.array([[c2, 0, s2], [0, 1, 0], [-s2, 0, c2]]) # around y 1171 1172 m3 = np.array([[1, 0, 0], [0, c3, -s3], [0, s3, c3]]) # around x 1173 1174 elif Y: 1175 m1 = np.array([[c1, -s1, 0], [s1, c1, 0], [0, 0, 1]]) # around z 1176 1177 m2 = np.array([[c2, 0, s2], [0, 1, 0], [-s2, 0, c2]]) # around y 1178 1179 m3 = np.array([[c3, -s3, 0], [s3, c3, 0], [0, 0, 1]]) # around z 1180 1181 else: 1182 m1 = np.array([[c1, -s1, 0], [s1, c1, 0], [0, 0, 1]]) # around z 1183 1184 m2 = np.array([[1, 0, 0], [0, c2, -s2], [0, s2, c2]]) # around x 1185 1186 m3 = np.array([[c3, -s3, 0], [s3, c3, 0], [0, 0, 1]]) # around z 1187 1188 M = np.dot(m3.T, np.dot(m2.T, m1.T)) 1189 1190 return M 1191 1192 1193def coordsys2euler_zyz(coord): 1194 """Return the ZYZ Euler angles corresponding to a rotation between 1195 the two coordinate systems. The angles can be used in rotate_alm. 1196 1197 Parameters 1198 ---------- 1199 coord : a tuple with initial and final coord systems. 1200 See normalise_coord. 1201 1202 Returns 1203 ------- 1204 psi, theta, phi : floats 1205 The Euler angles defining a ZYZ rotation between the coordinate 1206 systems. 1207 1208 Examples 1209 -------- 1210 >>> np.array(coordsys2euler_zyz('GE')) 1211 array([ 1.45937485, 1.05047962, -3.14119347]) 1212 >>> np.array(coordsys2euler_zyz('CG')) 1213 array([-0.22443941, 1.09730866, 2.14556934]) 1214 >>> np.array(coordsys2euler_zyz('E')) 1215 array([ 0., 0., 0.]) 1216 """ 1217 1218 coord_norm = normalise_coord(coord) 1219 1220 if coord_norm[0] == coord_norm[1]: 1221 psi, theta, phi = np.zeros(3) 1222 else: 1223 # Convert basis vectors 1224 1225 xin, yin, zin = np.eye(3) 1226 rot = Rotator(coord=coord_norm) 1227 xout, yout, zout = rot([xin, yin, zin]).T 1228 1229 # Normalize 1230 1231 xout /= np.sqrt(np.dot(xout, xout)) 1232 yout /= np.sqrt(np.dot(yout, yout)) 1233 zout /= np.sqrt(np.dot(zout, zout)) 1234 1235 # Get the angles 1236 1237 psi = np.arctan2(yout[2], -xout[2]) 1238 theta = np.arccos(np.dot(zin, zout)) 1239 phi = np.arctan2(zout[1], zout[0]) 1240 1241 return psi, theta, phi 1242