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