Merge pull request #1281 from lpsinger/iterative_all_world2pix

Add iterative wcs.all_world2pix

astropy/wcs/tests/test_wcs.py

@@ -5,15 +5,22 @@
 
 import numpy as np
 from numpy.testing import (
-    assert_array_almost_equal, assert_array_almost_equal_nulp)
+    assert_allclose, assert_array_almost_equal, assert_array_almost_equal_nulp)
 
-from ...tests.helper import raises, catch_warnings
+from ...tests.helper import raises, catch_warnings, pytest
 from ... import wcs
 from ...utils.data import (
     get_pkg_data_filenames, get_pkg_data_contents, get_pkg_data_filename)
 from ...tests.helper import pytest
 
 
+try:
+    import scipy
+except ImportError:
+    HAS_SCIPY = False
+else:
+    HAS_SCIPY = True
+
 # test_maps() is a generator
 def test_maps():
 
@@ -415,3 +422,24 @@ def test_validate():
     with open(get_pkg_data_filename("data/validate.txt"), "r") as fd:
         assert set([x.strip() for x in fd.readlines()]) == set([
             x.strip() for x in results_txt.splitlines()])
+
+
+@pytest.mark.skipif('not HAS_SCIPY')
+def test_all_world2pix():
+    """Test all_world2pix, iterative inverse of all_pix2world"""
+    fits = get_pkg_data_filename('data/sip.fits')
+    w = wcs.WCS(fits)
+
+    tolerance = 1e-6
+    world = 0.1 * np.random.randn(100, 2)
+    for i in range(len(w.wcs.crval)):
+        world[:, i] += w.wcs.crval[i]
+    all_pix = w.all_world2pix(world, 0, tolerance=tolerance)
+    wcs_pix = w.wcs_world2pix(world, 0)
+    all_world = w.all_pix2world(all_pix, 0)
+
+    # First, check that the SIP distortion correction at least produces
+    # some different answers from the WCS-only transform.
+    assert np.any(all_pix != wcs_pix)
+
+    assert_allclose(all_world, world, rtol=0, atol=tolerance)

astropy/wcs/wcs.py

@@ -1074,6 +1074,11 @@ def wcs_pix2world(self, *args, **kwargs):
 
         {1}
 
+        tolerance : float, optional
+            Tolerance of solution. Iteration terminates when the iterative
+            solver estimates that the true solution is within this many pixels
+            current estimate. Default value is 1e-6 (pixels).
+
         Returns
         -------
 
@@ -1122,6 +1127,94 @@ def wcs_pix2world(self, *args, **kwargs):
     def wcs_pix2sky(self, *args, **kwargs):
         return self.wcs_pix2world(*args, **kwargs)
 
+    def _all_world2pix(self, world, origin, tolerance, **kwargs):
+        try:
+            import scipy.optimize
+        except ImportError:
+            raise ImportError(
+                "You must have Scipy installed to use this method. " +
+                "See <http://www.scipy.org>.")
+        pix = []
+        for i in range(len(world)):
+            x0 = self.wcs_world2pix(np.atleast_2d(world[i]), origin,
+                **kwargs).flatten()
+            func = lambda pix: (self.all_pix2world(np.atleast_2d(pix),
+                origin, **kwargs) - world[i]).flatten()
+            # Use Broyden inverse because it is (a) present in a wide range of
+            # Scipy version, (b) provides an option for the absolute tolerance,
+            # and (c) is suitable for small-scale problems (i.e., a few
+            # variables, rather than hundreds of variables).
+            soln = scipy.optimize.broyden1(func, x0, x_tol=tolerance)
+            pix.append(soln.flatten())
+        return np.asarray(pix)
+
+    def all_world2pix(self, *args, **kwargs):
+        if self.wcs is None:
+            raise ValueError("No basic WCS settings were created.")
+        tolerance = kwargs.pop('tolerance', 1e-6)
+        return self._array_converter(lambda *args, **kwargs:
+            self._all_world2pix(*args, tolerance=tolerance, **kwargs),
+            'input', *args,
+            **kwargs)
+    all_world2pix.__doc__ = """
+        Transforms world coordinates to pixel coordinates, using numerical
+        iteration to invert the method `~astropy.wcs.WCS.all_pix2world` within a
+        tolerance of 1e-6 pixels.
+
+        Note that to use this function, you must have Scipy installed.
+
+        Parameters
+        ----------
+        {0}
+
+            For a transformation that is not two-dimensional, the
+            two-argument form must be used.
+
+        {1}
+
+        Returns
+        -------
+
+        {2}
+
+        Notes
+        -----
+        The order of the axes for the input world array is determined by
+        the `CTYPEia` keywords in the FITS header, therefore it may
+        not always be of the form (*ra*, *dec*).  The
+        `~astropy.wcs.Wcsprm.lat`, `~astropy.wcs.Wcsprm.lng`,
+        `~astropy.wcs.Wcsprm.lattyp` and `~astropy.wcs.Wcsprm.lngtyp`
+        members can be used to determine the order of the axes.
+
+        Raises
+        ------
+        MemoryError
+            Memory allocation failed.
+
+        SingularMatrixError
+            Linear transformation matrix is singular.
+
+        InconsistentAxisTypesError
+            Inconsistent or unrecognized coordinate axis types.
+
+        ValueError
+            Invalid parameter value.
+
+        ValueError
+            Invalid coordinate transformation parameters.
+
+        ValueError
+            x- and y-coordinate arrays are not the same size.
+
+        InvalidTransformError
+            Invalid coordinate transformation parameters.
+
+        InvalidTransformError
+            Ill-conditioned coordinate transformation parameters.
+        """.format(__.TWO_OR_MORE_ARGS('naxis', 8),
+                   __.RA_DEC_ORDER(8),
+                   __.RETURNS('pixel coordinates', 8))
+
     def wcs_world2pix(self, *args, **kwargs):
         if self.wcs is None:
             raise ValueError("No basic WCS settings were created.")