DM-26485: Add skyToPixelArray and pixelToSkyArray convenience functions.

python/lsst/afw/geom/skyWcs/skyWcsContinued.py

@@ -20,9 +20,79 @@
 # see <http://www.lsstcorp.org/LegalNotices/>.
 #
 
-__all__ = []
+import numpy as np
 
+from lsst.utils import continueClass
 from ..python import reduceTransform
 from .skyWcs import SkyWcs
 
+__all__ = []
+
+
+@continueClass  # noqa: F811
+class SkyWcs:
+    def pixelToSkyArray(self, x, y, degrees=False):
+        """
+        Convert numpy array pixels (x, y) to numpy array sky (ra, dec)
+        positions.
+
+        Parameters
+        ----------
+        x : `np.ndarray`
+            Array of x values.
+        y : `np.ndarray`
+            Array of y values.
+        degrees : `bool`, optional
+            Return ra, dec arrays in degrees if True.
+
+        Returns
+        -------
+        ra : `np.ndarray`
+            Array of Right Ascension.  Units are radians unless
+            degrees=True.
+        dec : `np.ndarray`
+            Array of Declination.  Units are radians unless
+            degrees=True.
+        """
+        xy = np.vstack((x, y))
+        ra, dec = np.vsplit(self.getTransform().getMapping().applyForward(xy), 2)
+        ra %= (2.*np.pi)
+
+        if degrees:
+            return np.rad2deg(ra.ravel()), np.rad2deg(dec.ravel())
+        else:
+            return ra.ravel(), dec.ravel()
+
+    def skyToPixelArray(self, ra, dec, degrees=False):
+        """
+        Convert numpy array sky (ra, dec) positions to numpy array
+        pixels (x, y).
+
+        Parameters
+        ----------
+        ra : `np.ndarray`
+            Array of Right Ascension.  Units are radians unless
+            degrees=True.
+        dec : `np.ndarray`
+            Array of Declination.  Units are radians unless
+            degrees=True.
+        degrees : `bool`, optional
+            Input ra, dec arrays are degrees if True.
+
+        Returns
+        -------
+        x : `np.ndarray`
+            Array of x values.
+        y : `np.ndarray`
+            Array of y values.
+        """
+        radec = np.vstack((ra, dec))
+        if degrees:
+            radec = np.deg2rad(radec)
+
+        x, y = np.vsplit(self.getTransform().getMapping().applyInverse(radec), 2)
+
+        return x.ravel(), y.ravel()
+
+
 SkyWcs.__reduce__ = reduceTransform

tests/test_skyWcs.py

@@ -7,6 +7,7 @@
 import astropy.coordinates
 import astropy.wcs
 import astshim as ast
+import numpy as np
 from numpy.testing import assert_allclose
 
 import lsst.utils.tests
@@ -545,6 +546,56 @@ def testAgainstAstropyWcs(self):
             # Most projections only seem to agree to within 1e-4 in the round trip test
             self.assertSkyWcsAstropyWcsAlmostEqual(skyWcs=skyWcs, astropyWcs=astropyWcs, bbox=bbox)
 
+    def testPixelToSkyArray(self):
+        """Test the numpy-array version of pixelToSky
+        """
+        cdMatrix = makeCdMatrix(scale=self.scale)
+        wcs = makeSkyWcs(crpix=self.crpix, crval=self.crvalList[0], cdMatrix=cdMatrix)
+
+        xPoints = np.array([0.0, 1000.0, 0.0, -1111.0])
+        yPoints = np.array([0.0, 0.0, 2000.0, -2222.0])
+
+        pixPointList = [lsst.geom.Point2D(x, y) for x, y in zip(xPoints, yPoints)]
+
+        spherePoints = wcs.pixelToSky(pixPointList)
+
+        ra, dec = wcs.pixelToSkyArray(xPoints, yPoints, degrees=False)
+        for r, d, spherePoint in zip(ra, dec, spherePoints):
+            self.assertAlmostEqual(r, spherePoint.getRa().asRadians())
+            self.assertAlmostEqual(d, spherePoint.getDec().asRadians())
+
+        ra, dec = wcs.pixelToSkyArray(xPoints, yPoints, degrees=True)
+        for r, d, spherePoint in zip(ra, dec, spherePoints):
+            self.assertAlmostEqual(r, spherePoint.getRa().asDegrees())
+            self.assertAlmostEqual(d, spherePoint.getDec().asDegrees())
+
+    def testSkyToPixelArray(self):
+        """Test the numpy-array version of skyToPixel
+        """
+        cdMatrix = makeCdMatrix(scale=self.scale)
+        wcs = makeSkyWcs(crpix=self.crpix, crval=self.crvalList[0], cdMatrix=cdMatrix)
+
+        raPoints = np.array([3.92646679e-02, 3.59646622e+02,
+                             3.96489283e-02, 4.70419353e-01])
+        decPoints = np.array([44.9722155, 44.97167735,
+                              45.52775599, 44.3540619])
+
+        spherePointList = [lsst.geom.SpherePoint(ra*lsst.geom.degrees,
+                                                 dec*lsst.geom.degrees)
+                           for ra, dec in zip(raPoints, decPoints)]
+
+        pixPoints = wcs.skyToPixel(spherePointList)
+
+        x, y = wcs.skyToPixelArray(np.deg2rad(raPoints), np.deg2rad(decPoints))
+        for x0, y0, pixPoint in zip(x, y, pixPoints):
+            self.assertAlmostEqual(x0, pixPoint.getX())
+            self.assertAlmostEqual(y0, pixPoint.getY())
+
+        x, y = wcs.skyToPixelArray(raPoints, decPoints, degrees=True)
+        for x0, y0, pixPoint in zip(x, y, pixPoints):
+            self.assertAlmostEqual(x0, pixPoint.getX())
+            self.assertAlmostEqual(y0, pixPoint.getY())
+
     def testStr(self):
         """Test that we can get something coherent when printing a SkyWcs.
         """