Compass drawing is a function

lsst-sitcom · Nov 1, 2023 · 9792ba3 · 9792ba3
1 parent 230c450
commit 9792ba3
Showing 1 changed file with 71 additions and 51 deletions.
diff --git a/python/lsst/summit/utils/plotting.py b/python/lsst/summit/utils/plotting.py
@@ -32,6 +32,76 @@
 import lsst.afw.image as afwImage
 
 
+def drawCompass(ax, wcs, compassLocation=300, arrowLength=300.):
+    """
+    Draw the compass.
+    The arrowLength is the length of compass arrows (arrows should have
+    the same length).
+    The steps here are:
+    - transform the (compassLocation, compassLocation) to RA, DEC coordinates
+    - move this point in DEC to get N; in RA to get E directions
+    - transform N and E points back to pixel coordinates
+    - find linear solutions for lines connecting the center of
+      the compass with N and E points
+    - find points along those lines located at the distance of
+      arrowLength form the (compassLocation, compassLocation).
+    - there will be two points for each linear solution.
+      Choose the correct one.
+    - centers of the N/E labels will also be located on those lines.
+
+    Parameters
+    ----------
+    ax : `matplotlib.axes.Axes`
+        The axes on which the compass will be drawn.
+    wcs : `lsst.afw.geom.SkyWcs`
+        WCS from exposure.
+    compassLocation : `int`, optional
+        How far in from the bottom left of the image to display the compass.
+    arrowLength : `float`, optional
+        The length of the compass arrow.
+    Returns
+    -------
+    ax : `matplotlib.axes.Axes`
+        The axes with the compass.
+    """
+
+    anchorRa, anchorDec = wcs.pixelToSky(compassLocation, compassLocation)
+    east = wcs.skyToPixel(geom.SpherePoint(anchorRa + 30.0 * geom.arcseconds, anchorDec))
+    north = wcs.skyToPixel(geom.SpherePoint(anchorRa, anchorDec + 30. * geom.arcseconds))
+    labelPosition = arrowLength + 50.
+
+    for xy, label in [(north, 'N'), (east, 'E')]:
+        if compassLocation == xy[0]:
+            xTip = compassLocation
+            xTipLabel = compassLocation
+            if xy[1] > compassLocation:
+                yTip = compassLocation + arrowLength
+                yTipLabel = compassLocation + labelPosition
+            else:
+                yTip = compassLocation - arrowLength
+                yTipLabel = compassLocation - labelPosition
+        else:
+            slope = (xy[1] - compassLocation) / (xy[0] - compassLocation)
+            xTipProjection = arrowLength / np.sqrt(1. + slope**2)
+            xTipLabelProjection = labelPosition / np.sqrt(1. + slope**2)
+
+            if xy[0] > compassLocation:
+                xTip = compassLocation + xTipProjection
+                xTipLabel = compassLocation + xTipLabelProjection
+            elif xy[0] < compassLocation:
+                xTip = compassLocation - xTipProjection
+                xTipLabel = compassLocation - xTipLabelProjection
+            yTip = slope * (xTip - compassLocation) + compassLocation
+            yTipLabel = slope * (xTipLabel - compassLocation) + compassLocation
+
+        color = 'r'
+        ax.arrow(compassLocation, compassLocation,
+                 xTip-compassLocation, yTip-compassLocation,
+                 head_width=30., length_includes_head=True, color=color)
+        ax.text(xTipLabel, yTipLabel, label, ha='center', va='center', color=color)
+    return ax
+
+
 def plot(inputData,
          figure=None,
          centroids=None,
@@ -169,58 +239,8 @@ def plot(inputData,
             wcs = None
 
         if wcs:
-            """
-            Place the compass.
-            The arrowLength is the length of compass arrows (arrows should have
-            the same length).
-            The steps here are:
-            - transform the (compassLocation, compassLocation)
-              to RA, DEC coordinates
-            - move this point in DEC to get N; in RA to get E directions
-            - transform N and E points back to pixel coordinates
-            - find linear solutions for lines connecting the center of
-              the compass with N and E points
-            - find points along those lines located at the distance of
-              arrowLength form the (compassLocation, compassLocation).
-            - there will be two points for each linear solution. Choose
-              the correct one.
-            - centers of the N/E labels will also be located on those lines.
-            """
-            anchorRa, anchorDec = wcs.pixelToSky(compassLocation, compassLocation)
-            east = wcs.skyToPixel(geom.SpherePoint(anchorRa + 30.0 * geom.arcseconds, anchorDec))
-            north = wcs.skyToPixel(geom.SpherePoint(anchorRa, anchorDec + 30. * geom.arcseconds))
             arrowLength = min(imageData.shape) * 0.05
-            labelPosition = arrowLength + 50.
-
-            for xy, label in [(north, 'N'), (east, 'E')]:
-                if compassLocation == xy[0]:
-                    xTip = compassLocation
-                    xTipLabel = compassLocation
-                    if xy[1] > compassLocation:
-                        yTip = compassLocation + arrowLength
-                        yTipLabel = compassLocation + labelPosition
-                    else:
-                        yTip = compassLocation - arrowLength
-                        yTipLabel = compassLocation - labelPosition
-                else:
-                    slope = (xy[1] - compassLocation) / (xy[0] - compassLocation)
-                    xTipProjection = arrowLength / np.sqrt(1. + slope**2)
-                    xTipLabelProjection = labelPosition / np.sqrt(1. + slope**2)
-
-                    if xy[0] > compassLocation:
-                        xTip = compassLocation + xTipProjection
-                        xTipLabel = compassLocation + xTipLabelProjection
-                    elif xy[0] < compassLocation:
-                        xTip = compassLocation - xTipProjection
-                        xTipLabel = compassLocation - xTipLabelProjection
-                    yTip = slope * (xTip - compassLocation) + compassLocation
-                    yTipLabel = slope * (xTipLabel - compassLocation) + compassLocation
-
-                color = 'r'
-                ax.arrow(compassLocation, compassLocation,
-                         xTip-compassLocation, yTip-compassLocation,
-                         head_width=30., length_includes_head=True, color=color)
-                ax.text(xTipLabel, yTipLabel, label, ha='center', va='center', color=color)
+            ax = drawCompass(ax, wcs, compassLocation=compassLocation, arrowLength=arrowLength)
 
     if centroids:
         ax.plot(*zip(*centroids),