Distortion Solution

corgidrp/data.py

@@ -1311,27 +1311,29 @@ class AstrometricCalibration(Image):
     Class for astrometric calibration file. 
 
     Args:
-        data_or_filepath (str or np.array); either the filepath to the FITS file to read in OR the 1D array of calibration measurements
+        data_or_filepath (str or np.array): either the filepath to the FITS file to read in OR an array of calibration measurements (boresight, plate scale, north angle, distortion coeffs)
         pri_hdr (astropy.io.fits.Header): the primary header (required only if raw 2D data is passed in)
         ext_hdr (astropy.io.fits.Header): the image extension header (required only if raw 2D data is passed in)
 
     Attrs:
         boresight (np.array): the [(RA, Dec)] of the center pixel in ([deg], [deg])
         platescale (float): the platescale value in [mas/pixel]
         northangle (float): the north angle value in [deg]
+        distortion_coeffs (np.array): the array of legendre polynomial coefficients that describe the distortion map (if distortion map is not computed this is an array of nans)
 
     """
-    def __init__(self, data_or_filepath, pri_hdr=None, ext_hdr=None, input_dataset=None):
+    def __init__(self, data_or_filepath, pri_hdr=None, ext_hdr=None, err=None, input_dataset=None):
         # run the image class constructor
-        super().__init__(data_or_filepath, pri_hdr=pri_hdr, ext_hdr=ext_hdr)
+        super().__init__(data_or_filepath, pri_hdr=pri_hdr, ext_hdr=ext_hdr, err=err)
 
         # File format checks
-        if self.data.shape != (4,):
-            raise ValueError("The AstrometricCalibration data should be a 1D array of four values")
+        if type(self.data) != np.ndarray:
+            raise ValueError("The AstrometricCalibration data should be an array of calibration measurements")
         else:
             self.boresight = self.data[:2]
             self.platescale = self.data[2]
             self.northangle = self.data[3]
+            self.distortion_coeffs = (self.data[4:-1], self.data[-1])
 
         # if this is a new astrometric calibration file, bookkeep it in the header
         # we need to check if it is new

corgidrp/mocks.py

@@ -866,17 +866,19 @@ def make_fluxmap_image(
         dq = dq)
     return image
 
-def create_astrom_data(field_path, filedir=None, subfield_radius=0.02, platescale=21.8, rotation=45, add_gauss_noise=True):
+def create_astrom_data(field_path, filedir=None, image_shape=(1024, 1024), subfield_radius=0.02, platescale=21.8, rotation=45, add_gauss_noise=True, distortion_coeffs_path=None):
     """
     Create simulated data for astrometric calibration.
 
     Args:
         field_path (str): Full path to directory with test field data (ra, dec, vmag, etc.)
         filedir (str): (Optional) Full path to directory to save to.
+        image_shape (tuple of ints): The desired shape of the image (num y pixels, num x pixels), (default: (1024, 1024))
         subfield_radius (float): The radius [deg] around the target coordinate for creating a subfield to produce the image from
         platescale (float): The plate scale of the created image data (default: 21.8 [mas/pixel])
         rotation (float): The north angle of the created image data (default: 45 [deg])
         add_gauss_noise (boolean): Argument to determine if gaussian noise should be added to the data (default: True)
+        distortion_coeffs_path (str): Full path to csv with the distortion coefficients and the order of polynomial used to describe distortion (default: None))
 
     Returns:
         corgidrp.data.Dataset:
@@ -891,9 +893,8 @@ def create_astrom_data(field_path, filedir=None, subfield_radius=0.02, platescal
         os.mkdir(filedir)
 
     # hard coded image properties
-    size = (1024, 1024)
-    sim_data = np.zeros(size)
-    ny, nx = size
+    sim_data = np.zeros(image_shape)
+    ny, nx = image_shape
     center = [nx //2, ny //2]
     target = (80.553428801, -69.514096821)
     fwhm = 3
@@ -936,6 +937,8 @@ def create_astrom_data(field_path, filedir=None, subfield_radius=0.02, platescal
 
     xpix = xpix[pix_inds]
     ypix = ypix[pix_inds]
+    ras = cal_SkyCoords[pix_inds]
+    decs = cal_SkyCoords[pix_inds]
 
     amplitudes = np.power(10, ((subfield['VMAG'][pix_inds] - 22.5) / (-2.5))) * 10  
 
@@ -984,9 +987,69 @@ def create_astrom_data(field_path, filedir=None, subfield_radius=0.02, platescal
         noise_rng = np.random.default_rng(10)
         gain = 1
         ref_flux = 10
-        noise = noise_rng.normal(scale= ref_flux/gain * 0.1, size= size)
+        noise = noise_rng.normal(scale= ref_flux/gain * 0.1, size= image_shape)
         sim_data = sim_data + noise
 
+    # add distortion (optional)
+    if distortion_coeffs_path is not None:
+        # load in distortion coeffs and fitorder
+        coeff_data = np.genfromtxt(distortion_coeffs_path, delimiter=',', dtype=None)
+        fitorder = int(coeff_data[-1])
+
+        # convert legendre polynomials into distortin maps in x and y 
+        yorig, xorig = np.indices(image_shape)
+        y0, x0 = image_shape[0]//2, image_shape[1]//2
+        yorig -= y0
+        xorig -= x0
+
+        # get the number of fitting params from the order
+        fitparams = (fitorder + 1)**2
+
+        # reshape the coeff arrays
+        best_params_x = coeff_data[:fitparams]
+        best_params_x = best_params_x.reshape(fitorder+1, fitorder+1)
+
+        total_orders = np.arange(fitorder+1)[:,None] + np.arange(fitorder+1)[None, :]
+
+        best_params_x = best_params_x / 500**(total_orders)
+
+        # evaluate the polynomial at all pixel positions
+        x_corr = np.polynomial.legendre.legval2d(xorig.ravel(), yorig.ravel(), best_params_x)
+        x_corr = x_corr.reshape(xorig.shape)
+        distmapx = x_corr - xorig
+
+        # reshape and evaluate the same for y
+        best_params_y = coeff_data[fitparams:-1]
+        best_params_y = best_params_y.reshape(fitorder+1, fitorder+1)
+
+        best_params_y = best_params_y / 500**(total_orders)
+
+        # evaluate the polynomial at all pixel positions
+        y_corr = np.polynomial.legendre.legval2d(xorig.ravel(), yorig.ravel(), best_params_y)
+        y_corr = y_corr.reshape(yorig.shape)
+        distmapy = y_corr - yorig
+
+        ## distort image based on coeffs
+        if (nx >= ny): imgsize = nx
+        else: imgsize = ny
+
+        gridx, gridy = np.meshgrid(np.arange(imgsize), np.arange(imgsize))
+        gridx = gridx + distmapx
+        gridy = gridy + distmapy
+
+        sim_data = scipy.ndimage.map_coordinates(sim_data, [gridy, gridx])
+
+        # transform the source coordinates
+        dist_xpix, dist_ypix = [], []
+        for (x,y) in zip(xpix, ypix):
+            x_new = x + distmapx[int(y)][int(x)]
+            y_new = y + distmapy[int(y)][int(x)]
+
+            dist_xpix.append(x_new)
+            dist_ypix.append(y_new)
+
+        xpix, ypix = dist_xpix, dist_ypix
+
     # load as an image object
     frames = []
     prihdr, exthdr = create_default_headers()
@@ -997,14 +1060,16 @@ def create_astrom_data(field_path, filedir=None, subfield_radius=0.02, platescal
     newhdr = fits.Header(new_hdr)
     frame = data.Image(sim_data, pri_hdr= prihdr, ext_hdr= newhdr)
     filename = "simcal_astrom.fits"
+    guessname = "guesses.csv"
     if filedir is not None:
         # save source SkyCoord locations and pixel location estimates
         guess = Table()
-        guess['x'] = [int(x) for x in xpix]
-        guess['y'] = [int(y) for y in ypix]
-        guess['RA'] = cal_SkyCoords[pix_inds].ra
-        guess['DEC'] = cal_SkyCoords[pix_inds].dec
-        ascii.write(guess, filedir+'/guesses.csv', overwrite=True)
+        guess['x'] = [x for x in xpix]
+        guess['y'] = [y for y in ypix]
+        guess['RA'] = ras.ra
+        guess['DEC'] = decs.dec
+        guess['VMAG'] = subfield['VMAG'][pix_inds]
+        ascii.write(guess, filedir+'/'+guessname, overwrite=True)
 
         frame.save(filedir=filedir, filename=filename)
 

tests/e2e_tests/astrom_e2e.py

@@ -4,7 +4,6 @@
 import numpy as np
 import astropy.time as time
 import astropy.io.fits as fits
-import astropy.io.ascii as ascii
 import corgidrp
 import corgidrp.data as data
 import corgidrp.mocks as mocks
@@ -178,7 +177,7 @@ def test_astrom_e2e(tvacdata_path, e2eoutput_path):
     assert dec == pytest.approx(target[1], abs=8.333e-7)
 
 if __name__ == "__main__":
-    tvacdata_dir = "/Users/macuser/Roman/corgi_contributions/Callibration_Notebooks/TVAC"
+    tvacdata_dir = "/Users/macuser/Roman/corgidrp_develop/calibration_notebooks/TVAC"
     outputdir = thisfile_dir
 
     ap = argparse.ArgumentParser(description="run the l1->l2b->boresight end-to-end test")

tests/test_astrom.py

@@ -31,7 +31,6 @@ def test_astrom():
 
     # perform the astrometric calibration
     astrom_cal = astrom.boresight_calibration(input_dataset=dataset, field_path=field_path, find_threshold=5)
-    assert len(astrom_cal.data) == 4
 
     # the data was generated to have the following image properties
     expected_platescale = 21.8

tests/test_data/distortion_expected_coeffs.csv

@@ -0,0 +1,33 @@
+0.00000000e+00  
+5.37588246e-01  
+1.45977937e-01 
+-5.71312982e-02
+5.00412506e+02  
+2.96600421e-01 
+-1.92137365e-01 
+-7.80792573e-01
+-4.85109411e-01 
+-2.75618370e-02 
+-1.58717081e-01 
+-1.94784434e-01
+-1.39807000e-02 
+-5.13642791e-02 
+-1.33704172e-02  
+1.08517713e-02
+0.00000000e+00 
+5.00855736e+02  
+3.15740375e-01 
+-1.46063718e-01
+-1.87760787e-02 
+-4.02194305e-01 
+-2.91035718e-02 
+1.72557478e-01
+-4.50792397e-01 
+-2.38286770e-01  
+1.81843117e-01  
+2.03185301e-01
+5.01096120e-01  
+2.54169286e-02  
+4.85195275e-02  
+5.61938400e-02
+3

tests/test_distortion_map_creation.py

@@ -0,0 +1,119 @@
+import os
+import pickle
+import numpy as np
+import pytest
+
+import corgidrp
+import corgidrp.mocks as mocks
+import corgidrp.astrom as astrom
+import corgidrp.data as data
+
+import astropy
+from astropy.io import ascii
+
+def test_distortion():
+    """ 
+    Generate a simulated image and test the distortion map creation as part of the boresight calibration.
+
+    """
+    # create a simulated image with source guesses and true positions
+    # check that the simulated image folder exists and create if not
+    datadir = os.path.join(os.path.dirname(__file__), "test_data", "simastrom")
+    if not os.path.exists(datadir):
+        os.mkdir(datadir)
+
+    field_path = os.path.join(os.path.dirname(__file__), "test_data", "JWST_CALFIELD2020.csv")
+    distortion_coeffs_path = os.path.join(os.path.dirname(__file__), "test_data", "distortion_expected_coeffs.csv")
+
+    mocks.create_astrom_data(field_path=field_path, filedir=datadir, distortion_coeffs_path=distortion_coeffs_path)
+
+    image_path = os.path.join(datadir, 'simcal_astrom.fits')
+    source_match_path = os.path.join(datadir, 'guesses.csv')
+    matches = ascii.read(source_match_path)
+
+    # open the image
+    dataset = data.Dataset([image_path])
+
+    # perform the astrometric calibration
+    # feed in the correct matches and use only ~100 randomly selected stars
+    rng = np.random.default_rng(seed=17)
+    select_stars = rng.choice(len(matches), size=150, replace=False)
+
+    astrom_cal = astrom.boresight_calibration(input_dataset=dataset, field_path=field_path, field_matches=[matches[select_stars]], find_distortion=True)
+
+    ## check that the distortion map does not create offsets greater than 4[mas]
+        # compute the distortion maps created from the best fit coeffs
+    coeffs = astrom_cal.distortion_coeffs
+    expected_coeffs = np.genfromtxt(distortion_coeffs_path)
+
+        # note the image shape and center around the image center
+    image_shape = np.shape(dataset[0].data)
+    yorig, xorig = np.indices(image_shape)
+    y0, x0 = image_shape[0]//2, image_shape[1]//2
+    yorig -= y0
+    xorig -= x0
+
+        # get the number of fitting params from the order
+    fitorder = int(coeffs[1])
+    fitparams = (fitorder + 1)**2
+    true_fitorder = int(expected_coeffs[-1])
+    true_fitparams = (true_fitorder + 1)**2
+
+        # reshape the coeff arrays
+    best_params_x = coeffs[0][:fitparams]
+    best_params_x = best_params_x.reshape(fitorder+1, fitorder+1)
+    total_orders = np.arange(fitorder+1)[:,None] + np.arange(fitorder+1)[None, :]
+    best_params_x = best_params_x / 500**(total_orders)
+
+    true_params_x = expected_coeffs[:-1][:true_fitparams]
+    true_params_x = true_params_x.reshape(true_fitorder+1, true_fitorder+1)
+    true_total_orders = np.arange(true_fitorder+1)[:,None] + np.arange(true_fitorder+1)[None, :]
+    true_params_x = true_params_x / 500**(true_total_orders)
+
+        # evaluate the polynomial at all pixel positions
+    x_corr = np.polynomial.legendre.legval2d(xorig.ravel(), yorig.ravel(), best_params_x)
+    x_corr = x_corr.reshape(xorig.shape)
+    x_diff = x_corr - xorig
+
+    true_x_corr = np.polynomial.legendre.legval2d(xorig.ravel(), yorig.ravel(), true_params_x)
+    true_x_corr = true_x_corr.reshape(xorig.shape)
+    true_x_diff = true_x_corr - xorig
+
+        # reshape and evaluate the same for y
+    best_params_y = coeffs[0][fitparams:]
+    best_params_y = best_params_y.reshape(fitorder+1, fitorder+1)
+    best_params_y = best_params_y / 500**(total_orders)
+
+    true_params_y = expected_coeffs[:-1][true_fitparams:]
+    true_params_y = true_params_y.reshape(true_fitorder+1, true_fitorder+1)
+    true_params_y = true_params_y / 500**(true_total_orders)
+
+        # evaluate the polynomial at all pixel positions
+    y_corr = np.polynomial.legendre.legval2d(xorig.ravel(), yorig.ravel(), best_params_y)
+    y_corr = y_corr.reshape(yorig.shape)
+    y_diff = y_corr - yorig
+
+    true_y_corr = np.polynomial.legendre.legval2d(xorig.ravel(), yorig.ravel(), true_params_y)
+    true_y_corr = true_y_corr.reshape(yorig.shape)
+    true_y_diff = true_y_corr - yorig
+
+    # check the distortion maps are less than the maximum injected distortion (~3 pixels)
+    assert np.all(np.abs(x_diff)) < np.max(np.abs(true_x_diff))
+    assert np.all(np.abs(y_diff)) < np.max(np.abs(true_y_diff))
+
+    # check they can be pickled (for CTC operations)
+    pickled = pickle.dumps(astrom_cal)
+    pickled_astrom = pickle.loads(pickled)
+    assert np.all((astrom_cal.data == pickled_astrom.data))
+
+    # save and check it can be pickled after save
+    astrom_cal.save(filedir=datadir, filename="astrom_cal_output.fits")
+    astrom_cal_2 = data.AstrometricCalibration(os.path.join(datadir, "astrom_cal_output.fits"))
+
+    # check they can be pickled (for CTC operations)
+    pickled = pickle.dumps(astrom_cal_2)
+    pickled_astrom = pickle.loads(pickled)
+    assert np.all((astrom_cal.data == pickled_astrom.data)) # check it is the same as the original
+
+if __name__ == "__main__":
+    test_distortion()