Merge pull request #432 from dirac-institute/fake_wcs

Consolidate fake WCS generation logic

src/kbmod/fake_data_creator.py

@@ -14,6 +14,7 @@
 from kbmod.configuration import SearchConfiguration
 from kbmod.file_utils import *
 from kbmod.search import *
+from kbmod.wcs_utils import append_wcs_to_hdu_header, make_fake_wcs_info
 from kbmod.work_unit import WorkUnit
 
 
@@ -204,22 +205,9 @@ def save_fake_data_to_dir(self, data_dir):
             img.save_layers(data_dir + "/")
 
             # Open the file and insert fake WCS data.
+            wcs_dict = make_fake_wcs_info(200.6145, -7.7888, 2000, 4000)
             hdul = fits.open(filename)
-            hdul[1].header["WCSAXES"] = 2
-            hdul[1].header["CTYPE1"] = "RA---TAN-SIP"
-            hdul[1].header["CTYPE2"] = "DEC--TAN-SIP"
-            hdul[1].header["CRVAL1"] = 200.614997245422
-            hdul[1].header["CRVAL2"] = -7.78878863332778
-            hdul[1].header["CRPIX1"] = 1033.934327
-            hdul[1].header["CRPIX2"] = 2043.548284
-            hdul[1].header["CD1_1"] = -1.13926485986789e-07
-            hdul[1].header["CD1_2"] = 7.31839748843125e-05
-            hdul[1].header["CD2_1"] = -7.30064978350695e-05
-            hdul[1].header["CD2_2"] = -1.27520156332774e-07
-            hdul[1].header["CTYPE1A"] = "LINEAR  "
-            hdul[1].header["CTYPE2A"] = "LINEAR  "
-            hdul[1].header["CUNIT1A"] = "PIXEL   "
-            hdul[1].header["CUNIT2A"] = "PIXEL   "
+            append_wcs_to_hdu_header(wcs_dict, hdul[1].header)
             hdul.writeto(filename, overwrite=True)
             hdul.close()
 

src/kbmod/wcs_utils.py

@@ -344,15 +344,19 @@ def append_wcs_to_hdu_header(wcs, header):
 
     Parameters
     ----------
-    wcs : `astropy.wcs.WCS`
-        The WCS to use.
+    wcs : `astropy.wcs.WCS` or `dict`
+        The WCS to use or a dictionary with the necessary information.
     header : `astropy.io.fits.Header`
         The header to which to append the data.
     """
     if wcs is not None:
-        wcs_header = wcs.to_header()
-        for key in wcs_header:
-            header[key] = wcs_header[key]
+        if type(wcs) is dict:
+            wcs_map = wcs
+        else:
+            wcs_map = wcs.to_header()
+
+        for key in wcs_map:
+            header[key] = wcs_map[key]
 
 
 def wcs_to_dict(wcs):
@@ -374,3 +378,46 @@ def wcs_to_dict(wcs):
         for key in wcs_header:
             result[key] = wcs_header[key]
     return result
+
+
+def make_fake_wcs_info(center_ra, center_dec, height, width, deg_per_pixel=None):
+    """Create a fake WCS given basic information. This is not a realistic
+    WCS in terms of astronomy, but can provide a place holder for many tests.
+
+    Parameters
+    ----------
+    center_ra : `float`
+        The center of the pointing in RA (degrees)
+    center_dev : `float`
+        The center of the pointing in DEC (degrees)
+    height : `int`
+        The image height (pixels).
+    width : `int`
+        The image width (pixels).
+    deg_per_pixel : `float`, optional
+        The angular resolution of a pixel (degrees).
+
+    Returns
+    -------
+    wcs_dict : `dict`
+        A dictionary with the information needed for a WCS.
+    """
+    wcs_dict = {
+        "WCSAXES": 2,
+        "CTYPE1": "RA---TAN-SIP",
+        "CTYPE2": "DEC--TAN-SIP",
+        "CRVAL1": center_ra,
+        "CRVAL2": center_dec,
+        "CRPIX1": height / 2.0,
+        "CRPIX2": width / 2.0,
+        "CTYPE1A": "LINEAR  ",
+        "CTYPE2A": "LINEAR  ",
+        "CUNIT1A": "PIXEL   ",
+        "CUNIT2A": "PIXEL   ",
+    }
+
+    if deg_per_pixel is not None:
+        wcs_dict["CDELT1"] = deg_per_pixel
+        wcs_dict["CDELT2"] = deg_per_pixel
+
+    return wcs_dict

tests/test_regression_test.py

@@ -17,6 +17,7 @@
 from kbmod.run_search import SearchRunner
 from kbmod.search import *
 from kbmod.trajectory_utils import make_trajectory
+from kbmod.wcs_utils import append_wcs_to_hdu_header, make_fake_wcs_info
 
 
 def ave_trajectory_distance(trjA, trjB, times=[0.0]):
@@ -205,21 +206,19 @@ def add_wcs_header_data(full_file_name):
         The path and filename of the FITS file to modify.
     """
     hdul = fits.open(full_file_name)
-    hdul[1].header["WCSAXES"] = 2
-    hdul[1].header["CTYPE1"] = "RA---TAN-SIP"
-    hdul[1].header["CTYPE2"] = "DEC--TAN-SIP"
-    hdul[1].header["CRVAL1"] = 200.614997245422
-    hdul[1].header["CRVAL2"] = -7.78878863332778
-    hdul[1].header["CRPIX1"] = 1033.934327
-    hdul[1].header["CRPIX2"] = 2043.548284
+
+    # Create a fake WCS with some minimal information.
+    wcs_dict = make_fake_wcs_info(200.615, -7.789, 2068, 4088)
+    append_wcs_to_hdu_header(wcs_dict, hdul[1].header)
+
+    # Add rotational and scale parameters that will allow us to use a matching
+    # set of search angles (using KBMOD angle suggestion function).
     hdul[1].header["CD1_1"] = -1.13926485986789e-07
     hdul[1].header["CD1_2"] = 7.31839748843125e-05
     hdul[1].header["CD2_1"] = -7.30064978350695e-05
     hdul[1].header["CD2_2"] = -1.27520156332774e-07
-    hdul[1].header["CTYPE1A"] = "LINEAR  "
-    hdul[1].header["CTYPE2A"] = "LINEAR  "
-    hdul[1].header["CUNIT1A"] = "PIXEL   "
-    hdul[1].header["CUNIT2A"] = "PIXEL   "
+
+    # Save the augmented header.
     hdul.writeto(full_file_name, overwrite=True)
     hdul.close()
 

tests/test_wcs_utils.py

@@ -45,16 +45,54 @@ def test_wcs_to_dict(self):
             self.assertAlmostEqual(new_dict[key], self.header_dict[key])
 
     def test_append_wcs_to_hdu_header(self):
-        pri = fits.PrimaryHDU()
-        self.assertFalse("CRVAL1" in pri.header)
-        self.assertFalse("CRVAL2" in pri.header)
-        self.assertFalse("CRPIX1" in pri.header)
-        self.assertFalse("CRPIX2" in pri.header)
-
-        append_wcs_to_hdu_header(self.wcs, pri.header)
-        for key in self.header_dict:
-            self.assertTrue(key in pri.header)
-            self.assertAlmostEqual(pri.header[key], self.header_dict[key])
+        for use_dictionary in [True, False]:
+            if use_dictionary:
+                wcs_info = self.header_dict
+            else:
+                wcs_info = self.wcs
+
+            # Run the dictionary and full WCS tests as separate subtests
+            with self.subTest(i=use_dictionary):
+                pri = fits.PrimaryHDU()
+                self.assertFalse("CRVAL1" in pri.header)
+                self.assertFalse("CRVAL2" in pri.header)
+                self.assertFalse("CRPIX1" in pri.header)
+                self.assertFalse("CRPIX2" in pri.header)
+
+                append_wcs_to_hdu_header(wcs_info, pri.header)
+                for key in self.header_dict:
+                    self.assertTrue(key in pri.header)
+                    self.assertAlmostEqual(pri.header[key], self.header_dict[key])
+
+    def test_make_fake_wcs_info(self):
+        # Test that we make the dictionary
+        wcs_dict = make_fake_wcs_info(25.0, -10.0, 200, 100, deg_per_pixel=0.01)
+        self.assertEqual(wcs_dict["WCSAXES"], 2)
+        self.assertEqual(wcs_dict["CTYPE1"], "RA---TAN-SIP")
+        self.assertEqual(wcs_dict["CTYPE2"], "DEC--TAN-SIP")
+        self.assertEqual(wcs_dict["CRVAL1"], 25.0)
+        self.assertEqual(wcs_dict["CRVAL2"], -10.0)
+        self.assertEqual(wcs_dict["CRPIX1"], 100.0)
+        self.assertEqual(wcs_dict["CRPIX2"], 50.0)
+        self.assertEqual(wcs_dict["CDELT1"], 0.01)
+        self.assertEqual(wcs_dict["CDELT2"], 0.01)
+        self.assertEqual(wcs_dict["CTYPE1A"], "LINEAR  ")
+        self.assertEqual(wcs_dict["CTYPE2A"], "LINEAR  ")
+        self.assertEqual(wcs_dict["CUNIT1A"], "PIXEL   ")
+        self.assertEqual(wcs_dict["CUNIT2A"], "PIXEL   ")
+
+        # Test that we can convert to a WCS and perform predictions.
+        test_wcs = WCS(wcs_dict)
+
+        # Center position should be at approximately (25.0, -10.0).
+        pos = test_wcs.pixel_to_world(99, 49)
+        self.assertAlmostEqual(pos.ra.degree, 25.0, delta=0.001)
+        self.assertAlmostEqual(pos.dec.degree, -10.0, delta=0.001)
+
+        # One pixel off position should be at approximately (25.01, -10.0).
+        pos = test_wcs.pixel_to_world(100, 48)
+        self.assertAlmostEqual(pos.ra.degree, 25.01, delta=0.01)
+        self.assertAlmostEqual(pos.dec.degree, -10.0, delta=0.01)
 
 
 class test_construct_wcs_tangent_projection(unittest.TestCase):

tests/test_work_unit.py

@@ -8,6 +8,7 @@
 import warnings
 
 from kbmod.configuration import SearchConfiguration
+from kbmod.fake_data_creator import make_fake_wcs_info
 import kbmod.search as kb
 from kbmod.work_unit import hdu_to_raw_image, raw_image_to_hdu, WorkUnit
 
@@ -44,23 +45,7 @@ def setUp(self):
         self.config.set("repeated_flag_keys", None)
 
         # Create a fake WCS
-        header_dict = {
-            "WCSAXES": 2,
-            "CTYPE1": "RA---TAN-SIP",
-            "CTYPE2": "DEC--TAN-SIP",
-            "CRVAL1": 200.614997245422,
-            "CRVAL2": -7.78878863332778,
-            "CRPIX1": 1033.934327,
-            "CRPIX2": 2043.548284,
-            "CD1_1": -1.13926485986789e-07,
-            "CD1_2": 7.31839748843125e-05,
-            "CD2_1": -7.30064978350695e-05,
-            "CD2_2": -1.27520156332774e-07,
-            "CTYPE1A": "LINEAR  ",
-            "CTYPE2A": "LINEAR  ",
-            "CUNIT1A": "PIXEL   ",
-            "CUNIT2A": "PIXEL   ",
-        }
+        header_dict = make_fake_wcs_info(200.6145, -7.7888, 2000, 4000)
         self.wcs = WCS(header_dict)
         self.per_image_wcs = per_image_wcs = [
             (self.wcs if i % 2 == 0 else None) for i in range(self.num_images)