1# -*- coding: utf-8 -*-
2r"""
3==========================================================
4Create a new coordinate class (for the Sagittarius stream)
5==========================================================
6
7This document describes in detail how to subclass and define a custom spherical
8coordinate frame, as discussed in :ref:`astropy:astropy-coordinates-design` and
9the docstring for `~astropy.coordinates.BaseCoordinateFrame`. In this example,
10we will define a coordinate system defined by the plane of orbit of the
11Sagittarius Dwarf Galaxy (hereafter Sgr; as defined in Majewski et al. 2003).
12The Sgr coordinate system is often referred to in terms of two angular
13coordinates, :math:`\Lambda,B`.
14
15To do this, we need to define a subclass of
16`~astropy.coordinates.BaseCoordinateFrame` that knows the names and units of the
17coordinate system angles in each of the supported representations.  In this case
18we support `~astropy.coordinates.SphericalRepresentation` with "Lambda" and
19"Beta". Then we have to define the transformation from this coordinate system to
20some other built-in system. Here we will use Galactic coordinates, represented
21by the `~astropy.coordinates.Galactic` class.
22
23See Also
24--------
25
26* The `gala package <http://gala.adrian.pw/>`_, which defines a number of
27  Astropy coordinate frames for stellar stream coordinate systems.
28* Majewski et al. 2003, "A Two Micron All Sky Survey View of the Sagittarius
29  Dwarf Galaxy. I. Morphology of the Sagittarius Core and Tidal Arms",
30  https://arxiv.org/abs/astro-ph/0304198
31* Law & Majewski 2010, "The Sagittarius Dwarf Galaxy: A Model for Evolution in a
32  Triaxial Milky Way Halo", https://arxiv.org/abs/1003.1132
33* David Law's Sgr info page https://www.stsci.edu/~dlaw/Sgr/
34
35
36*By: Adrian Price-Whelan, Erik Tollerud*
37
38*License: BSD*
39
40
41"""
42
43##############################################################################
44# Make `print` work the same in all versions of Python, set up numpy,
45# matplotlib, and use a nicer set of plot parameters:
46
47import numpy as np
48import matplotlib.pyplot as plt
49from astropy.visualization import astropy_mpl_style
50plt.style.use(astropy_mpl_style)
51
52
53##############################################################################
54# Import the packages necessary for coordinates
55
56from astropy.coordinates import frame_transform_graph
57from astropy.coordinates.matrix_utilities import rotation_matrix, matrix_product, matrix_transpose
58import astropy.coordinates as coord
59import astropy.units as u
60
61##############################################################################
62# The first step is to create a new class, which we'll call
63# ``Sagittarius`` and make it a subclass of
64# `~astropy.coordinates.BaseCoordinateFrame`:
65
66
67class Sagittarius(coord.BaseCoordinateFrame):
68    """
69    A Heliocentric spherical coordinate system defined by the orbit
70    of the Sagittarius dwarf galaxy, as described in
71        https://ui.adsabs.harvard.edu/abs/2003ApJ...599.1082M
72    and further explained in
73        https://www.stsci.edu/~dlaw/Sgr/.
74
75    Parameters
76    ----------
77    representation : `~astropy.coordinates.BaseRepresentation` or None
78        A representation object or None to have no data (or use the other keywords)
79    Lambda : `~astropy.coordinates.Angle`, optional, must be keyword
80        The longitude-like angle corresponding to Sagittarius' orbit.
81    Beta : `~astropy.coordinates.Angle`, optional, must be keyword
82        The latitude-like angle corresponding to Sagittarius' orbit.
83    distance : `~astropy.units.Quantity`, optional, must be keyword
84        The Distance for this object along the line-of-sight.
85    pm_Lambda_cosBeta : `~astropy.units.Quantity`, optional, must be keyword
86        The proper motion along the stream in ``Lambda`` (including the
87        ``cos(Beta)`` factor) for this object (``pm_Beta`` must also be given).
88    pm_Beta : `~astropy.units.Quantity`, optional, must be keyword
89        The proper motion in Declination for this object (``pm_ra_cosdec`` must
90        also be given).
91    radial_velocity : `~astropy.units.Quantity`, optional, keyword-only
92        The radial velocity of this object.
93
94    """
95
96    default_representation = coord.SphericalRepresentation
97    default_differential = coord.SphericalCosLatDifferential
98
99    frame_specific_representation_info = {
100        coord.SphericalRepresentation: [
101            coord.RepresentationMapping('lon', 'Lambda'),
102            coord.RepresentationMapping('lat', 'Beta'),
103            coord.RepresentationMapping('distance', 'distance')]
104    }
105
106
107##############################################################################
108# Breaking this down line-by-line, we define the class as a subclass of
109# `~astropy.coordinates.BaseCoordinateFrame`. Then we include a descriptive
110# docstring.  The final lines are class-level attributes that specify the
111# default representation for the data, default differential for the velocity
112# information, and mappings from the attribute names used by representation
113# objects to the names that are to be used by the ``Sagittarius`` frame. In this
114# case we override the names in the spherical representations but don't do
115# anything with other representations like cartesian or cylindrical.
116#
117# Next we have to define the transformation from this coordinate system to some
118# other built-in coordinate system; we will use Galactic coordinates. We can do
119# this by defining functions that return transformation matrices, or by simply
120# defining a function that accepts a coordinate and returns a new coordinate in
121# the new system. Because the transformation to the Sagittarius coordinate
122# system is just a spherical rotation from Galactic coordinates, we'll just
123# define a function that returns this matrix. We'll start by constructing the
124# transformation matrix using pre-determined Euler angles and the
125# ``rotation_matrix`` helper function:
126
127SGR_PHI = (180 + 3.75) * u.degree # Euler angles (from Law & Majewski 2010)
128SGR_THETA = (90 - 13.46) * u.degree
129SGR_PSI = (180 + 14.111534) * u.degree
130
131# Generate the rotation matrix using the x-convention (see Goldstein)
132D = rotation_matrix(SGR_PHI, "z")
133C = rotation_matrix(SGR_THETA, "x")
134B = rotation_matrix(SGR_PSI, "z")
135A = np.diag([1.,1.,-1.])
136SGR_MATRIX = matrix_product(A, B, C, D)
137
138
139##############################################################################
140# Since we already constructed the transformation (rotation) matrix above, and
141# the inverse of a rotation matrix is just its transpose, the required
142# transformation functions are very simple:
143
144@frame_transform_graph.transform(coord.StaticMatrixTransform, coord.Galactic, Sagittarius)
145def galactic_to_sgr():
146    """ Compute the transformation matrix from Galactic spherical to
147        heliocentric Sgr coordinates.
148    """
149    return SGR_MATRIX
150
151
152##############################################################################
153# The decorator ``@frame_transform_graph.transform(coord.StaticMatrixTransform,
154# coord.Galactic, Sagittarius)``  registers this function on the
155# ``frame_transform_graph`` as a coordinate transformation. Inside the function,
156# we simply return the previously defined rotation matrix.
157#
158# We then register the inverse transformation by using the transpose of the
159# rotation matrix (which is faster to compute than the inverse):
160
161@frame_transform_graph.transform(coord.StaticMatrixTransform, Sagittarius, coord.Galactic)
162def sgr_to_galactic():
163    """ Compute the transformation matrix from heliocentric Sgr coordinates to
164        spherical Galactic.
165    """
166    return matrix_transpose(SGR_MATRIX)
167
168
169##############################################################################
170# Now that we've registered these transformations between ``Sagittarius`` and
171# `~astropy.coordinates.Galactic`, we can transform between *any* coordinate
172# system and ``Sagittarius`` (as long as the other system has a path to
173# transform to `~astropy.coordinates.Galactic`). For example, to transform from
174# ICRS coordinates to ``Sagittarius``, we would do:
175
176icrs = coord.SkyCoord(280.161732*u.degree, 11.91934*u.degree, frame='icrs')
177sgr = icrs.transform_to(Sagittarius)
178print(sgr)
179
180##############################################################################
181# Or, to transform from the ``Sagittarius`` frame to ICRS coordinates (in this
182# case, a line along the ``Sagittarius`` x-y plane):
183
184sgr = coord.SkyCoord(Lambda=np.linspace(0, 2*np.pi, 128)*u.radian,
185                     Beta=np.zeros(128)*u.radian, frame='sagittarius')
186icrs = sgr.transform_to(coord.ICRS)
187print(icrs)
188
189##############################################################################
190# As an example, we'll now plot the points in both coordinate systems:
191
192fig, axes = plt.subplots(2, 1, figsize=(8, 10),
193                         subplot_kw={'projection': 'aitoff'})
194
195axes[0].set_title("Sagittarius")
196axes[0].plot(sgr.Lambda.wrap_at(180*u.deg).radian, sgr.Beta.radian,
197             linestyle='none', marker='.')
198
199axes[1].set_title("ICRS")
200axes[1].plot(icrs.ra.wrap_at(180*u.deg).radian, icrs.dec.radian,
201             linestyle='none', marker='.')
202
203plt.show()
204
205##############################################################################
206# This particular transformation is just a spherical rotation, which is a
207# special case of an Affine transformation with no vector offset. The
208# transformation of velocity components is therefore natively supported as
209# well:
210
211sgr = coord.SkyCoord(Lambda=np.linspace(0, 2*np.pi, 128)*u.radian,
212                     Beta=np.zeros(128)*u.radian,
213                     pm_Lambda_cosBeta=np.random.uniform(-5, 5, 128)*u.mas/u.yr,
214                     pm_Beta=np.zeros(128)*u.mas/u.yr,
215                     frame='sagittarius')
216icrs = sgr.transform_to(coord.ICRS)
217print(icrs)
218
219fig, axes = plt.subplots(3, 1, figsize=(8, 10), sharex=True)
220
221axes[0].set_title("Sagittarius")
222axes[0].plot(sgr.Lambda.degree,
223             sgr.pm_Lambda_cosBeta.value,
224             linestyle='none', marker='.')
225axes[0].set_xlabel(r"$\Lambda$ [deg]")
226axes[0].set_ylabel(
227    fr"$\mu_\Lambda \, \cos B$ [{sgr.pm_Lambda_cosBeta.unit.to_string('latex_inline')}]")
228
229axes[1].set_title("ICRS")
230axes[1].plot(icrs.ra.degree, icrs.pm_ra_cosdec.value,
231             linestyle='none', marker='.')
232axes[1].set_ylabel(
233    fr"$\mu_\alpha \, \cos\delta$ [{icrs.pm_ra_cosdec.unit.to_string('latex_inline')}]")
234
235axes[2].set_title("ICRS")
236axes[2].plot(icrs.ra.degree, icrs.pm_dec.value,
237             linestyle='none', marker='.')
238axes[2].set_xlabel("RA [deg]")
239axes[2].set_ylabel(
240    fr"$\mu_\delta$ [{icrs.pm_dec.unit.to_string('latex_inline')}]")
241
242plt.show()
243