{
  "nbformat_minor": 0, 
  "nbformat": 4, 
  "cells": [
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "%matplotlib inline"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }, 
    {
      "source": [
        "\n==========================================================\nCreate a new coordinate class (for the Sagittarius stream)\n==========================================================\n\nThis document describes in detail how to subclass and define a custom spherical\ncoordinate frame, as discussed in :ref:`astropy-coordinates-design` and the\ndocstring for `~astropy.coordinates.BaseCoordinateFrame`. In this example, we\nwill define a coordinate system defined by the plane of orbit of the Sagittarius\nDwarf Galaxy (hereafter Sgr; as defined in Majewski et al. 2003).  The Sgr\ncoordinate system is often referred to in terms of two angular coordinates,\n$\\Lambda,B$.\n\nTo do this, wee need to define a subclass of\n`~astropy.coordinates.BaseCoordinateFrame` that knows the names and units of the\ncoordinate system angles in each of the supported representations.  In this case\nwe support `~astropy.coordinates.SphericalRepresentation` with \"Lambda\" and\n\"Beta\". Then we have to define the transformation from this coordinate system to\nsome other built-in system. Here we will use Galactic coordinates, represented\nby the `~astropy.coordinates.Galactic` class.\n\nSee Also\n--------\n\n* Majewski et al. 2003, \"A Two Micron All Sky Survey View of the Sagittarius\n  Dwarf Galaxy. I. Morphology of the Sagittarius Core and Tidal Arms\",\n  http://arxiv.org/abs/astro-ph/0304198\n* Law & Majewski 2010, \"The Sagittarius Dwarf Galaxy: A Model for Evolution in a\n  Triaxial Milky Way Halo\", http://arxiv.org/abs/1003.1132\n* David Law's Sgr info page http://www.astro.virginia.edu/~srm4n/Sgr/\n\n-------------------\n\n*By: Adrian Price-Whelan, Erik Tollerud*\n\n*License: BSD*\n\n-------------------\n\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "source": [
        "Make `print` work the same in all versions of Python, set up numpy,\nmatplotlib, and use a nicer set of plot parameters:\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "from __future__ import print_function\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom astropy.visualization import astropy_mpl_style\nplt.style.use(astropy_mpl_style)"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }, 
    {
      "source": [
        "Import the packages necessary for coordinates\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "from astropy.coordinates import frame_transform_graph\nfrom astropy.coordinates.angles import rotation_matrix\nimport astropy.coordinates as coord\nimport astropy.units as u"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }, 
    {
      "source": [
        "The first step is to create a new class, which we'll call\n``Sagittarius`` and make it a subclass of\n`~astropy.coordinates.BaseCoordinateFrame`:\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "class Sagittarius(coord.BaseCoordinateFrame):\n    \"\"\"\n    A Heliocentric spherical coordinate system defined by the orbit\n    of the Sagittarius dwarf galaxy, as described in\n        http://adsabs.harvard.edu/abs/2003ApJ...599.1082M\n    and further explained in\n        http://www.astro.virginia.edu/~srm4n/Sgr/.\n\n    Parameters\n    ----------\n    representation : `BaseRepresentation` or None\n        A representation object or None to have no data (or use the other keywords)\n    Lambda : `Angle`, optional, must be keyword\n        The longitude-like angle corresponding to Sagittarius' orbit.\n    Beta : `Angle`, optional, must be keyword\n        The latitude-like angle corresponding to Sagittarius' orbit.\n    distance : `Quantity`, optional, must be keyword\n        The Distance for this object along the line-of-sight.\n\n    \"\"\"\n    default_representation = coord.SphericalRepresentation\n\n    frame_specific_representation_info = {\n        'spherical': [coord.RepresentationMapping('lon', 'Lambda'),\n                      coord.RepresentationMapping('lat', 'Beta'),\n                      coord.RepresentationMapping('distance', 'distance')],\n        'unitspherical': [coord.RepresentationMapping('lon', 'Lambda'),\n                          coord.RepresentationMapping('lat', 'Beta')]\n    }"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }, 
    {
      "source": [
        "Breaking this down line-by-line, we define the class as a subclass of\n`~astropy.coordinates.BaseCoordinateFrame`. Then we include a descriptive\ndocstring.  The final lines are class-level attributes that specify the\ndefault representation for the data and mappings from the attribute names used\nby representation objects to the names that are to be used by ``Sagittarius``.\nIn this case we override the names in the spherical representations but don't\ndo anything with other representations like cartesian or cylindrical.\n\nNext we have to define the transformation from this coordinate system to some\nother built-in coordinate system; we will use Galactic coordinates. We can do\nthis by defining functions that return transformation matrices, or by simply\ndefining a function that accepts a coordinate and returns a new coordinate in\nthe new system. We'll start by constructing the rotation matrix, using the\n``rotation_matrix`` helper function:\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "SGR_PHI = np.radians(180+3.75) # Euler angles (from Law & Majewski 2010)\nSGR_THETA = np.radians(90-13.46)\nSGR_PSI = np.radians(180+14.111534)\n\n# Generate the rotation matrix using the x-convention (see Goldstein)\nD = rotation_matrix(SGR_PHI, \"z\", unit=u.radian)\nC = rotation_matrix(SGR_THETA, \"x\", unit=u.radian)\nB = rotation_matrix(SGR_PSI, \"z\", unit=u.radian)\nSGR_MATRIX = np.array(B.dot(C).dot(D))"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }, 
    {
      "source": [
        "This is done at the module level, since it will be used by both the\ntransformation from Sgr to Galactic as well as the inverse from Galactic to\nSgr. Now we can define our first transformation function:\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "@frame_transform_graph.transform(coord.FunctionTransform, coord.Galactic, Sagittarius)\ndef galactic_to_sgr(gal_coord, sgr_frame):\n    \"\"\" Compute the transformation from Galactic spherical to\n        heliocentric Sgr coordinates.\n    \"\"\"\n\n    l = np.atleast_1d(gal_coord.l.radian)\n    b = np.atleast_1d(gal_coord.b.radian)\n\n    X = np.cos(b)*np.cos(l)\n    Y = np.cos(b)*np.sin(l)\n    Z = np.sin(b)\n\n    # Calculate X,Y,Z,distance in the Sgr system\n    Xs, Ys, Zs = SGR_MATRIX.dot(np.array([X, Y, Z]))\n    Zs = -Zs\n\n    # Calculate the angular coordinates lambda,beta\n    Lambda = np.arctan2(Ys,Xs)*u.radian\n    Lambda[Lambda < 0] = Lambda[Lambda < 0] + 2.*np.pi*u.radian\n    Beta = np.arcsin(Zs/np.sqrt(Xs*Xs+Ys*Ys+Zs*Zs))*u.radian\n\n    return Sagittarius(Lambda=Lambda, Beta=Beta,\n                       distance=gal_coord.distance)"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }, 
    {
      "source": [
        "The decorator ``@frame_transform_graph.transform(coord.FunctionTransform,\ncoord.Galactic, Sagittarius)``  registers this function on the\n``frame_transform_graph`` as a coordinate transformation. Inside the function,\nwe simply follow the same procedure as detailed by David Law's `transformation\ncode <http://www.astro.virginia.edu/~srm4n/Sgr/code.html>`_. Note that in this\ncase, both coordinate systems are heliocentric, so we can simply copy any\ndistance from the `~astropy.coordinates.Galactic` object.\n\nWe then register the inverse transformation by using the transpose of the\nrotation matrix (which is faster to compute than the inverse):\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "@frame_transform_graph.transform(coord.FunctionTransform, Sagittarius, coord.Galactic)\ndef sgr_to_galactic(sgr_coord, gal_frame):\n    \"\"\" Compute the transformation from heliocentric Sgr coordinates to\n        spherical Galactic.\n    \"\"\"\n    L = np.atleast_1d(sgr_coord.Lambda.radian)\n    B = np.atleast_1d(sgr_coord.Beta.radian)\n\n    Xs = np.cos(B)*np.cos(L)\n    Ys = np.cos(B)*np.sin(L)\n    Zs = np.sin(B)\n    Zs = -Zs\n\n    X, Y, Z = SGR_MATRIX.T.dot(np.array([Xs, Ys, Zs]))\n\n    l = np.arctan2(Y,X)*u.radian\n    b = np.arcsin(Z/np.sqrt(X*X+Y*Y+Z*Z))*u.radian\n\n    l[l<=0] += 2*np.pi*u.radian\n\n    return coord.Galactic(l=l, b=b, distance=sgr_coord.distance)"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }, 
    {
      "source": [
        "Now that we've registered these transformations between ``Sagittarius`` and\n`~astropy.coordinates.Galactic`, we can transform between *any* coordinate\nsystem and ``Sagittarius`` (as long as the other system has a path to\ntransform to `~astropy.coordinates.Galactic`). For example, to transform from\nICRS coordinates to ``Sagittarius``, we simply:\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "icrs = coord.ICRS(280.161732*u.degree, 11.91934*u.degree)\nsgr = icrs.transform_to(Sagittarius)\nprint(sgr)"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }, 
    {
      "source": [
        "Or, to transform from the ``Sagittarius`` frame to ICRS coordinates:\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "sgr = Sagittarius(Lambda=np.linspace(0,2*np.pi,128)*u.radian,\n                  Beta=np.zeros(128)*u.radian)\nicrs = sgr.transform_to(coord.ICRS)"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }, 
    {
      "source": [
        "Plot the points in both coordinate systems:\n"
      ], 
      "cell_type": "markdown", 
      "metadata": {}
    }, 
    {
      "execution_count": null, 
      "cell_type": "code", 
      "source": [
        "fig,axes = plt.subplots(2, 1, figsize=(6,12),\n                        subplot_kw={'projection': 'aitoff'})\n\naxes[0].set_title(\"Sagittarius\")\naxes[0].plot(sgr.Lambda.wrap_at(180*u.deg).radian, sgr.Beta.radian, linestyle='none')\n\naxes[1].set_title(\"ICRS\")\naxes[1].plot(icrs.ra.wrap_at(180*u.deg).radian, icrs.dec.radian, linestyle='none')\n\nplt.show()"
      ], 
      "outputs": [], 
      "metadata": {
        "collapsed": false
      }
    }
  ], 
  "metadata": {
    "kernelspec": {
      "display_name": "Python 2", 
      "name": "python2", 
      "language": "python"
    }, 
    "language_info": {
      "mimetype": "text/x-python", 
      "nbconvert_exporter": "python", 
      "name": "python", 
      "file_extension": ".py", 
      "version": "2.7.11", 
      "pygments_lexer": "ipython2", 
      "codemirror_mode": {
        "version": 2, 
        "name": "ipython"
      }
    }
  }
}