create fit_barycentric_wcs (#541)

* add fit_barycentric_wcs

* add dimmension check

* remove scipy import

* fix indents

src/kbmod/reprojection_utils.py

@@ -1,7 +1,8 @@
 import astropy.units as u
 import numpy as np
 from astropy import units as u
-from astropy.coordinates import GCRS, ICRS
+from astropy.coordinates import SkyCoord, GCRS, ICRS
+from astropy.wcs.utils import fit_wcs_from_points
 from scipy.optimize import minimize
 
 
@@ -55,3 +56,54 @@ def correct_parallax(coord, obstime, point_on_earth, guess_distance):
     ).transform_to(ICRS())
 
     return answer
+
+
+def fit_barycentric_wcs(original_wcs, width, height, distance, obstime, point_on_earth, npoints=10):
+    """Given a ICRS WCS and an object's distance from the Sun,
+    return a new WCS that has been corrected for parallax motion.
+
+    Attributes
+    ----------
+    original_wcs : `astropy.wcs.WCS`
+        The image's WCS.
+    width : `int`
+        The image's width (typically NAXIS1).
+    height : `int`
+        The image's height (typically NAXIS2).
+    distance : `float`
+        The distance of the object from the sun, in AU.
+    obstime : `astropy.time.Time` or `string`
+        The observation time.
+    point_on_earth : `astropy.coordinate.EarthLocation`
+        The location on Earth of the observation.
+    npoints : `int`
+        The number of randomly sampled points to use during the WCS fitting.
+        Typically, the more points the higher the accuracy. The four corners
+        of the image will always be included, so setting npoints = 0 will mean
+        just using the corners.
+
+    Returns
+    ----------
+    An `astropy.wcs.WCS` representing the original image in "Explicity Barycentric Distance" (EBD)
+    space, i.e. where the points have been corrected for parallax.
+    """
+    sampled_x_points = np.array([0, 0, width, width])
+    sampled_y_points = np.array([0, height, height, 0])
+    if npoints > 0:
+        sampled_x_points = np.append(sampled_x_points, np.random.rand(npoints) * width)
+        sampled_y_points = np.append(sampled_y_points, np.random.rand(npoints) * height)
+
+    sampled_ra, sampled_dec = original_wcs.all_pix2world(sampled_x_points, sampled_y_points, 0)
+
+    sampled_coordinates = SkyCoord(sampled_ra, sampled_dec, unit="deg")
+
+    ebd_corrected_points = []
+    for coord in sampled_coordinates:
+        ebd_corrected_points.append(correct_parallax(coord, obstime, point_on_earth, distance))
+
+    ebd_corrected_points = SkyCoord(ebd_corrected_points)
+    xy = (sampled_x_points, sampled_y_points)
+    ebd_wcs = fit_wcs_from_points(
+        xy, ebd_corrected_points, proj_point="center", projection="TAN", sip_degree=3
+    )
+    return ebd_wcs

tests/test_reprojection_utils.py

@@ -1,10 +1,12 @@
 import unittest
 
+import numpy as np
 import numpy.testing as npt
 from astropy.coordinates import EarthLocation, SkyCoord, solar_system_ephemeris
 from astropy.time import Time
+from astropy.wcs import WCS
 
-from kbmod.reprojection_utils import correct_parallax
+from kbmod.reprojection_utils import correct_parallax, fit_barycentric_wcs
 
 
 class test_reprojection_utils(unittest.TestCase):
@@ -45,3 +47,71 @@ def test_parallax_equinox(self):
 
         npt.assert_almost_equal(corrected_coord2.ra.value, expected_ra)
         npt.assert_almost_equal(corrected_coord2.dec.value, expected_dec)
+
+    def test_fit_barycentric_wcs(self):
+        nx = 2046
+        ny = 4094
+        test_wcs = WCS(naxis=2)
+        test_wcs.pixel_shape = (ny, nx)
+        test_wcs.wcs.crpix = [nx / 2, ny / 2]
+        test_wcs.wcs.cdelt = np.array([-0.000055555555556, 0.000055555555556])
+        test_wcs.wcs.crval = [346.9681342111, -6.482196848597]
+        test_wcs.wcs.ctype = ["RA---TAN", "DEC--TAN"]
+
+        x_points = np.array([247, 1252, 1052, 980, 420, 1954, 730, 1409, 1491, 803])
+
+        y_points = np.array([1530, 713, 3414, 3955, 1975, 123, 1456, 2008, 1413, 1756])
+
+        expected_ra = np.array(
+            [
+                346.69225567,
+                346.63734563,
+                346.64836252,
+                346.65231188,
+                346.68282256,
+                346.59898412,
+                346.66587788,
+                346.62881986,
+                346.6243199,
+                346.66190162,
+            ]
+        )
+
+        expected_dec = np.array(
+            [
+                -6.62151717,
+                -6.66580019,
+                -6.51929901,
+                -6.48995635,
+                -6.5973741,
+                -6.6977762,
+                -6.62551611,
+                -6.59555108,
+                -6.62782211,
+                -6.60924105,
+            ]
+        )
+
+        expected_sc = SkyCoord(ra=expected_ra, dec=expected_dec, unit="deg")
+
+        time = "2021-08-24T20:59:06"
+        site = "ctio"
+        loc = EarthLocation.of_site(site)
+        distance = 41.1592725489203
+
+        corrected_wcs = fit_barycentric_wcs(
+            test_wcs,
+            nx,
+            ny,
+            distance,
+            time,
+            loc,
+        )
+
+        corrected_ra, corrected_dec = corrected_wcs.all_pix2world(x_points, y_points, 0)
+        corrected_sc = SkyCoord(corrected_ra, corrected_dec, unit="deg")
+        seps = expected_sc.separation(corrected_sc).arcsecond
+
+        # assert we have sub-milliarcsecond precision
+        assert np.all(seps < 0.001)
+        assert corrected_wcs.array_shape == (ny, nx)