diff --git a/src/kbmod/fake_data_creator.py b/src/kbmod/fake_data_creator.py index 7f15b5a4..86e07c7d 100644 --- a/src/kbmod/fake_data_creator.py +++ b/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() diff --git a/src/kbmod/wcs_utils.py b/src/kbmod/wcs_utils.py index 0a59d7ab..2c8687c8 100644 --- a/src/kbmod/wcs_utils.py +++ b/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 diff --git a/tests/test_regression_test.py b/tests/test_regression_test.py index 5b8ce73b..3794b17f 100644 --- a/tests/test_regression_test.py +++ b/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() diff --git a/tests/test_wcs_utils.py b/tests/test_wcs_utils.py index 0d0434a3..628a701c 100644 --- a/tests/test_wcs_utils.py +++ b/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 the we can convert to a WCS and extra 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): diff --git a/tests/test_work_unit.py b/tests/test_work_unit.py index b1c4b7fd..0e1c2f32 100644 --- a/tests/test_work_unit.py +++ b/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)