1# Licensed under a 3-clause BSD style license - see LICENSE.rst
2
3
4import numpy as np
5import warnings
6from matplotlib.patches import Polygon
7
8from astropy import units as u
9from astropy.coordinates import SkyCoord
10from astropy.coordinates.representation import UnitSphericalRepresentation, SphericalRepresentation
11from astropy.coordinates.matrix_utilities import rotation_matrix, matrix_product
12from astropy.utils.exceptions import AstropyUserWarning
13
14
15__all__ = ['Quadrangle', 'SphericalCircle']
16
17# Monkey-patch the docs to fix CapStyle and JoinStyle subs.
18# TODO! delete when upstream fix matplotlib/matplotlib#19839
19Polygon.__init__.__doc__ = Polygon.__init__.__doc__.replace(
20    "`.CapStyle`", "``matplotlib._enums.CapStyle``")
21Polygon.__init__.__doc__ = Polygon.__init__.__doc__.replace(
22    "`.JoinStyle`", "``matplotlib._enums.JoinStyle``")
23Polygon.set_capstyle.__doc__ = Polygon.set_capstyle.__doc__.replace(
24    "`.CapStyle`", "``matplotlib._enums.CapStyle``")
25Polygon.set_joinstyle.__doc__ = Polygon.set_joinstyle.__doc__.replace(
26    "`.JoinStyle`", "``matplotlib._enums.JoinStyle``")
27
28
29def _rotate_polygon(lon, lat, lon0, lat0):
30    """
31    Given a polygon with vertices defined by (lon, lat), rotate the polygon
32    such that the North pole of the spherical coordinates is now at (lon0,
33    lat0). Therefore, to end up with a polygon centered on (lon0, lat0), the
34    polygon should initially be drawn around the North pole.
35    """
36
37    # Create a representation object
38    polygon = UnitSphericalRepresentation(lon=lon, lat=lat)
39
40    # Determine rotation matrix to make it so that the circle is centered
41    # on the correct longitude/latitude.
42    m1 = rotation_matrix(-(0.5 * np.pi * u.radian - lat0), axis='y')
43    m2 = rotation_matrix(-lon0, axis='z')
44    transform_matrix = matrix_product(m2, m1)
45
46    # Apply 3D rotation
47    polygon = polygon.to_cartesian()
48    polygon = polygon.transform(transform_matrix)
49    polygon = UnitSphericalRepresentation.from_cartesian(polygon)
50
51    return polygon.lon, polygon.lat
52
53
54class SphericalCircle(Polygon):
55    """
56    Create a patch representing a spherical circle - that is, a circle that is
57    formed of all the points that are within a certain angle of the central
58    coordinates on a sphere. Here we assume that latitude goes from -90 to +90
59
60    This class is needed in cases where the user wants to add a circular patch
61    to a celestial image, since otherwise the circle will be distorted, because
62    a fixed interval in longitude corresponds to a different angle on the sky
63    depending on the latitude.
64
65    Parameters
66    ----------
67    center : tuple or `~astropy.units.Quantity` ['angle']
68        This can be either a tuple of two `~astropy.units.Quantity` objects, or
69        a single `~astropy.units.Quantity` array with two elements
70        or a `~astropy.coordinates.SkyCoord` object.
71    radius : `~astropy.units.Quantity` ['angle']
72        The radius of the circle
73    resolution : int, optional
74        The number of points that make up the circle - increase this to get a
75        smoother circle.
76    vertex_unit : `~astropy.units.Unit`
77        The units in which the resulting polygon should be defined - this
78        should match the unit that the transformation (e.g. the WCS
79        transformation) expects as input.
80
81    Notes
82    -----
83    Additional keyword arguments are passed to `~matplotlib.patches.Polygon`
84    """
85
86    def __init__(self, center, radius, resolution=100, vertex_unit=u.degree, **kwargs):
87
88        # Extract longitude/latitude, either from a SkyCoord object, or
89        # from a tuple of two quantities or a single 2-element Quantity.
90        # The SkyCoord is converted to SphericalRepresentation, if not already.
91        if isinstance(center, SkyCoord):
92            rep_type = center.representation_type
93            if not issubclass(rep_type, (SphericalRepresentation,
94                                         UnitSphericalRepresentation)):
95                warnings.warn(f'Received `center` of representation type {rep_type} '
96                              'will be converted to SphericalRepresentation ',
97                              AstropyUserWarning)
98            longitude, latitude = center.spherical.lon, center.spherical.lat
99        else:
100            longitude, latitude = center
101
102        # Start off by generating the circle around the North pole
103        lon = np.linspace(0., 2 * np.pi, resolution + 1)[:-1] * u.radian
104        lat = np.repeat(0.5 * np.pi - radius.to_value(u.radian), resolution) * u.radian
105
106        lon, lat = _rotate_polygon(lon, lat, longitude, latitude)
107
108        # Extract new longitude/latitude in the requested units
109        lon = lon.to_value(vertex_unit)
110        lat = lat.to_value(vertex_unit)
111
112        # Create polygon vertices
113        vertices = np.array([lon, lat]).transpose()
114
115        super().__init__(vertices, **kwargs)
116
117
118class Quadrangle(Polygon):
119    """
120    Create a patch representing a latitude-longitude quadrangle.
121
122    The edges of the quadrangle lie on two lines of constant longitude and two
123    lines of constant latitude (or the equivalent component names in the
124    coordinate frame of interest, such as right ascension and declination).
125    Note that lines of constant latitude are not great circles.
126
127    Unlike `matplotlib.patches.Rectangle`, the edges of this patch will render
128    as curved lines if appropriate for the WCS transformation.
129
130    Parameters
131    ----------
132    anchor : tuple or `~astropy.units.Quantity` ['angle']
133        This can be either a tuple of two `~astropy.units.Quantity` objects, or
134        a single `~astropy.units.Quantity` array with two elements.
135    width : `~astropy.units.Quantity` ['angle']
136        The width of the quadrangle in longitude (or, e.g., right ascension)
137    height : `~astropy.units.Quantity` ['angle']
138        The height of the quadrangle in latitude (or, e.g., declination)
139    resolution : int, optional
140        The number of points that make up each side of the quadrangle -
141        increase this to get a smoother quadrangle.
142    vertex_unit : `~astropy.units.Unit` ['angle']
143        The units in which the resulting polygon should be defined - this
144        should match the unit that the transformation (e.g. the WCS
145        transformation) expects as input.
146
147    Notes
148    -----
149    Additional keyword arguments are passed to `~matplotlib.patches.Polygon`
150    """
151
152    def __init__(self, anchor, width, height, resolution=100, vertex_unit=u.degree, **kwargs):
153
154        # Extract longitude/latitude, either from a tuple of two quantities, or
155        # a single 2-element Quantity.
156        longitude, latitude = u.Quantity(anchor).to_value(vertex_unit)
157
158        # Convert the quadrangle dimensions to the appropriate units
159        width = width.to_value(vertex_unit)
160        height = height.to_value(vertex_unit)
161
162        # Create progressions in longitude and latitude
163        lon_seq = longitude + np.linspace(0, width, resolution + 1)
164        lat_seq = latitude + np.linspace(0, height, resolution + 1)
165
166        # Trace the path of the quadrangle
167        lon = np.concatenate([lon_seq[:-1],
168                              np.repeat(lon_seq[-1], resolution),
169                              np.flip(lon_seq[1:]),
170                              np.repeat(lon_seq[0], resolution)])
171        lat = np.concatenate([np.repeat(lat_seq[0], resolution),
172                              lat_seq[:-1],
173                              np.repeat(lat_seq[-1], resolution),
174                              np.flip(lat_seq[1:])])
175
176        # Create polygon vertices
177        vertices = np.array([lon, lat]).transpose()
178
179        super().__init__(vertices, **kwargs)
180