diff --git a/corgidrp/astrom.py b/corgidrp/astrom.py index 46ae0252..714262d5 100644 --- a/corgidrp/astrom.py +++ b/corgidrp/astrom.py @@ -73,7 +73,7 @@ def measure_offset(frame, xstar_guess, ystar_guess, xoffset_guess, yoffset_guess Returns: binary_offset (np.array): List of [x,y] offsets in respective directions - + fit_errs (np.array): Array of [x,y] fitting errors """ #### Centroid on location of star ### yind = int(ystar_guess) @@ -125,10 +125,11 @@ def measure_offset(frame, xstar_guess, ystar_guess, xoffset_guess, yoffset_guess ### Fit the PSF to the data ### popt, pcov = optimize.curve_fit(shift_psf, stamp, data.ravel(), p0=(0,0,guessflux), maxfev=2000) tinyoffsets = popt[0:2] + fit_errs = np.sqrt([pcov[0,0], pcov[1,1], pcov[2,2]]) binary_offset = [xoffset_guess - tinyoffsets[0], yoffset_guess - tinyoffsets[1]] - return binary_offset + return binary_offset, fit_errs def compute_combinations(iteration, r=2): """ @@ -368,7 +369,9 @@ def match_sources(image, sources, field_path, comparison_threshold=50, rad=0.007 initial_platescale = np.mean(np.array([best_l1 / l1, best_l2 / l2, best_l3 / l3])) # [deg/mas] # find pseudo north angle from difference in triangle rotations from the target value - rot_image = np.array([angle_between((image.ext_hdr['CRPIX1'], image.ext_hdr['CRPIX2']), (s['x'], s['y'])) for s in [source1, source2, source3]]) + # assume target star is at the center of the image + ymid, xmid = image.data.shape + rot_image = np.array([angle_between((xmid //2, ymid //2), (s['x'], s['y'])) for s in [source1, source2, source3]]) rot_field = np.array([target_skycoord.position_angle(t).deg for t in skycoords[[best_sky_ind]]]) initial_northangle = np.abs(np.mean(rot_field - rot_image)) @@ -437,6 +440,86 @@ def match_sources(image, sources, field_path, comparison_threshold=50, rad=0.007 return matched_image_to_field +def fit_distortion_solution(params, fitorder, platescale, rotangle, pos1, meas_offset, sky_offset, meas_errs): + ''' + Cost function used to fit the legendre polynomials for distortion mapping. + + Args: + params (list): List of the x and y legendre polynomial coefficients + fitorder (int): The degree of legendre polynomial being used + platescale (float): The platescale of the image + rotangle (float): The north angle of the image + pos1 (np.array): A (2 x N) array of (x, y) pixel positions for the first star in N pairs + meas_offset (np.array): A (2 x N) array of (x, y) pixel offset from the first star position for N star pairs + sky_offset (np.array): A (2 x N) array of (sep, pa) true sky offsets in [mas] and [deg] from the first star position for N pairs + meas_errs (np.array): A (2 x N) array of (x, y) pixel errors in measured offsets from the first star position for N pairs + x0 (int): The x pixel value for the center of the image/ detector + y0 (int): The y pixel value for the center of the image/ detector + + Returns: + residuals (list): List of residuals between true and measured star positions + ''' + + fitparams = (fitorder + 1)**2 + + leg_params_x = np.array(params[:fitparams]) # the first half of params are for x fitting + leg_params_x = leg_params_x.reshape(fitorder+1, fitorder+1) + + leg_params_y = np.array(params[fitparams:]) # the last half are for y fitting + leg_params_y = leg_params_y.reshape(fitorder+1, fitorder+1) + + total_orders = np.arange(fitorder+1)[:,None] + np.arange(fitorder+1)[None,:] # creating a 4 x 4 matrix of order numbers (?) + + leg_params_x = leg_params_x / 500**(total_orders) # making the coefficients sufficiently large for fitting (or else ~0) + leg_params_y = leg_params_y / 500**(total_orders) + + residuals = [] + + binary_offsets = meas_offset.T + star1_pos = pos1.T + + # derive star2 position + star2_pos = star1_pos + binary_offsets + + # undistort x and y for both star positions + star1_pos_corr = np.copy(star1_pos) + star1_pos_corr[:,0] = np.polynomial.legendre.legval2d(star1_pos[:,0], star1_pos[:,1], leg_params_x) + star1_pos_corr[:,1] = np.polynomial.legendre.legval2d(star1_pos[:,0], star1_pos[:,1], leg_params_y) + star2_pos_corr = np.copy(star2_pos) + star2_pos_corr[:,0] = np.polynomial.legendre.legval2d(star2_pos[:,0], star2_pos[:,1], leg_params_x) + star2_pos_corr[:,1] = np.polynomial.legendre.legval2d(star2_pos[:,0], star2_pos[:,1], leg_params_y) + + # translate offsets from [mas] sep, pa to [pixel] sep, pa + sky_sep = sky_offset[0] + sky_pa = sky_offset[1] + + true_offset_sep = sky_sep / platescale + true_offset_pa = sky_pa - rotangle + + # translate star_pos_corr from x, y to sep, pa + corr_offset_x = star2_pos_corr[:,0] - star1_pos_corr[:,0] + corr_offset_y = star2_pos_corr[:,1] - star1_pos_corr[:,1] + + corr_offset_sep = np.sqrt(corr_offset_x**2 + corr_offset_y**2) + corr_offset_pa = np.degrees(np.arctan2(-corr_offset_x, corr_offset_y)) + + res_sep = corr_offset_sep - true_offset_sep # this is in pixels + + res_pa_arctan_num = np.sin(np.radians(corr_offset_pa - true_offset_pa)) + res_pa_arctan_denom = np.cos(np.radians(corr_offset_pa - true_offset_pa)) + res_pa = np.arctan(res_pa_arctan_num / res_pa_arctan_denom) # this is in radians + + # translate pixel error in measurement to sep pa errs + sep_err = np.sqrt(meas_errs[0]**2 + meas_errs[1]**2) + pa_err = np.arctan2(-meas_errs[0], meas_errs[1]) # this is in radians + + res_sep /= sep_err + res_pa /= pa_err + + residuals = np.append(residuals, np.array([res_sep, res_pa]).ravel()) + + return residuals + def compute_platescale_and_northangle(image, source_info, center_coord, center_radius=0.9): """ Used to find the platescale and north angle of the image. Calculates the platescale for each pair of stars in the image @@ -507,7 +590,7 @@ def compute_platescale_and_northangle(image, source_info, center_coord, center_r xguess = star2['x'] - star1['x'] yguess = star2['y'] - star1['y'] - xoff, yoff = measure_offset(image, xstar_guess=star1['x'], ystar_guess=star2['x'], xoffset_guess= xguess, yoffset_guess= yguess) + (xoff, yoff), _ = measure_offset(image, xstar_guess=star1['x'], ystar_guess=star1['y'], xoffset_guess= xguess, yoffset_guess= yguess) pixsep = np.sqrt(np.power(xoff,2) + np.power(yoff,2)) pixseps[i] = pixsep @@ -649,14 +732,130 @@ def compute_boresight(image, source_info, target_coordinate, cal_properties): return image_center_RA, image_center_DEC -def boresight_calibration(input_dataset, field_path='JWST_CALFIELD2020.csv', field_matches=None, find_threshold=10, fwhm=7, mask_rad=1, comparison_threshold=50, search_rad=0.0075, platescale_guess=21.8, platescale_tol=0.1, center_radius=0.9, frames_to_combine=None): +def format_distortion_inputs(input_dataset, source_matches, ref_star_pos, position_error=None): + ''' Function that formats the input data for the distortion map computation * must be run before compute_distortion * + + Args: + input_dataset (corgidrp.data.dataset): corgidrp dataset object with images to compute the distortion from + source_matches (list of astropy.table.Table() objects): List of length N for N frames in the input dataset. Tables must columns 'x','y','RA','DEC' as pixel locations and corresponding sky positons + ref_star_pos (list of astropy.table.Table() objects): List of length N for N frames. Tables must have column names 'x', 'y', 'RA', 'DEC' for the position of the reference star to compute pairs with + position_error (NoneType or int): If int, this is the uniform error value assumed for the offset between pairs of stars in both x and y + Should be changed later to accept non-uniform errors + + Returns: + first_stars (np.array): 2D array of the (x, y) pixel positions for the first star in every star pair + offsets (np.array): 2D array of the (delta_x, delta_y) values for each star from the first star position + true_offsets (np.array): 2D array of the (delta_ra, delta_dec) offsets between the matched stars in the reference field + errs (np.array): 2D array of the (x_err, y_err) error in the measured pixel positions + ''' + ## Create arrays to store values in + dxs, dys = np.array([]), np.array([]) + xerrs, yerrs = np.array([]), np.array([]) + firstxs, firstys = np.array([]), np.array([]) + seps, pas = np.array([]), np.array([]) + + ## Loop over every frame in the dataset + for frame_ind in range(len(input_dataset)): + input_image = input_dataset[frame_ind].data + + # create all combinations of the target star with all others + combo_list = range(len(source_matches[frame_ind])) + skycoords = SkyCoord(ra= source_matches[frame_ind]['RA'], dec= source_matches[frame_ind]['DEC'], unit='deg', frame='icrs') + target_coord = SkyCoord(ra= ref_star_pos[frame_ind]['RA'], dec= ref_star_pos[frame_ind]['DEC'], unit='deg', frame='icrs') + + for pair_ind in combo_list: + # get the pixel offset + star1 = ref_star_pos[frame_ind] + star2 = source_matches[frame_ind][pair_ind] + + x_guess = star2['x'] - star1['x'] + y_guess = star2['y'] - star1['y'] + + (dx, dy), (xfit_err, yfit_err, _) = measure_offset(input_image, star1['x'], star1['y'], x_guess, y_guess, guessflux=10000) + + # get the true sky offset [mas] + true1 = target_coord + true2 = skycoords[pair_ind] + + # get true sky separation and position angle + true_sep = true1.separation(true2).mas + true_pa = true1.position_angle(true2).deg + + dxs = np.append(dxs, dx) + dys = np.append(dys, dy) + firstxs = np.append(firstxs, star1['x']) + firstys = np.append(firstys, star1['y']) + seps = np.append(seps, true_sep) + pas = np.append(pas, true_pa) + + if type(position_error) == type(None): + xerrs = np.append(xerrs, xfit_err) + yerrs = np.append(yerrs, yfit_err) + else: + xerrs = np.append(xerrs, position_error) + yerrs = np.append(yerrs, position_error) + + # join arrays and reshape to (1, 2, N) + offsets = np.array([dxs, dys]) + first_stars = np.array([firstxs, firstys]) + true_offsets = np.array([seps, pas]) + errs = np.array([xerrs, yerrs]) + + return first_stars, offsets, true_offsets, errs + +def compute_distortion(input_dataset, pos1, meas_offset, sky_offset, meas_errs, platescale, northangle, fitorder=3, initial_guess=None): + ''' + Function that computes the legendre polynomial coefficients that describe the image distortion map * must run format_disotrtio_inputs() first * + + Args: + input_dataset (corgidrp.data.Dataset): corgidrp dataset object with images to compute the distortion from + pos1 (np.array): 2D array of the (x, y) pixel positions for the first star in every star pair + meas_offset (np.array): 2D array of the (delta_x, delta_y) values for each star from the first star position + sky_offset (np.array): 2D array of the (delta_ra, delta_dec) offsets between the matched stars in the reference field + meas_errs (np.array): 2D array of the (x_err, y_err) error in the measured pixel positions + platescale (float): Platescale value to use in computing distortion + northangle (float): Northangle value to use in computing distortion + fitorder (int): The order of legendre polynomial to fit to the image distortion (default: 3) + initial_guess (np.array): Initial guess of fitting parameters (legendre coefficients) length based on fitorder (2 * (fitorder+1)**2), (default: None) + + Returns: + distortion_coeffs (tuple): The legendre coefficients (np.array) and polynomial order used for the fit (int) + + ''' + + ## SET FITTING PARAMS + # assume all images in dataset have the same shape + input_image = input_dataset[0].data + x0 = np.shape(input_image)[1] // 2 + y0 = np.shape(input_image)[0] // 2 + + # define fitting params + fitparams = (fitorder + 1)**2 + + # initial guesses for the legendre coeffs if none are passed + if initial_guess is None: + initial_guess = [0 for _ in range(fitorder+1)] + [500,] + [0 for _ in range(fitparams - fitorder - 2)] + [0,500] + [0 for _ in range(fitparams-2)] + + ## OPTIMIZE + # first_stars_, offsets_, true_offsets_, errs_ = first_stars, offsets, true_offsets, errs + (distortion_coeffs, _) = optimize.leastsq(fit_distortion_solution, initial_guess, + args=(fitorder, platescale, + northangle, pos1, meas_offset, + sky_offset, meas_errs)) + + return (distortion_coeffs, fitorder) + + +def boresight_calibration(input_dataset, field_path='JWST_CALFIELD2020.csv', field_matches=None, find_threshold=10, fwhm=7, mask_rad=1, + comparison_threshold=50, search_rad=0.0075, platescale_guess=21.8, platescale_tol=0.1, center_radius=0.9, + frames_to_combine=None, find_distortion=False, fitorder=3, position_error=None, initial_dist_guess=None): """ Perform the boresight calibration of a dataset. Args: input_dataset (corgidrp.data.Dataset): Dataset containing a images for astrometric calibration field_path (str): Full path to file with search field data (ra, dec, vmag, etc.) (default: 'JWST_CALFIELD2020.csv') - field_matches (str): Full path to file with calibration field matches to the image sources (x, y, ra, dec), if None, automated source matching is used (default: None) + field_matches (list of str or astropy.table.Table): List of full paths to files or astropy tables with calibration field matches for each image in the dataset (x, y, ra, dec), if single str the same filepath used for all frames,nif None, automated source matching is used (default: None) find_threshold (int): Number of stars to find (default 10) fwhm (float): Full width at half maximum of the stellar psf (default: 7, ~fwhm for a normal distribution with sigma=3) mask_rad (int): Radius of mask for stars [in fwhm] (default: 1) @@ -666,6 +865,10 @@ def boresight_calibration(input_dataset, field_path='JWST_CALFIELD2020.csv', fie platescale_tol (float): A tolerance for finding source matches within a fraction of the initial plate scale guess (default: 0.1) center_radius (float): Percent of the image to compute plate scale and north angle from, centered around the image center (default: 0.9 -- ie: 90% of the image is used) frames_to_combine (int): The number of frames to combine in a dataset (default: None) + find_distortion (boolean): Used to determine if distortion map coeffs will be computed (default: False) + fitorder (int): The order of legendre polynomials used to fit the distortion map (default: 3) + position_error (NoneType or int): If int, this is the uniform error value assumed for the offset between pairs of stars in both x and y + initial_dist_guess (np.array): An initial guess of legendre coefficients used for fitting distortion, if None will use coeffs associated with no distortion (default: None) Returns: corgidrp.data.AstrometricCalibration: Astrometric Calibration data object containing image center coords in (RA,DEC), platescale, and north angle @@ -675,8 +878,28 @@ def boresight_calibration(input_dataset, field_path='JWST_CALFIELD2020.csv', fie dataset = input_dataset.copy() # load in the source matches if automated source finder is not being used + matched_sources_multiframe = [] if field_matches is not None: - matched_sources = ascii.read(field_matches) + if len(field_matches) == 1: # single str or astropy.table case + for i in range(len(dataset)): + if type(field_matches[0]) == str: + matched_sources = ascii.read(field_matches[0]) + matched_sources_multiframe.append(matched_sources) + else: + matched_sources_multiframe.append(field_matches[0]) + # elif len(field_matches) == len(dataset): # this needs to be if the len(field_matches >1) + elif len(field_matches) > 1: # unique matches for each frame case + if len(field_matches) != len(dataset): + raise TypeError('field_matches must be a single str/ astropy.table OR the same length as input_dataset') + else: + for i in range(len(field_matches)): + if type(field_matches[0]) == str: + matched_sources = ascii.read(field_matches[i]) + matched_sources_multiframe.append(matched_sources) + else: + matched_sources_multiframe = field_matches + # else: + # raise TypeError('field_matches must be a single str or the same length as input_dataset') # load in field data to refer to if field_path == 'JWST_CALFIELD2020.csv': @@ -707,6 +930,7 @@ def boresight_calibration(input_dataset, field_path='JWST_CALFIELD2020.csv', fie # create a place to store all the calibration measurements astroms = [] + target_coord_tables = [] for i in range(len(dataset)): in_dataset = corgidrp.data.Dataset([dataset[i]]) @@ -714,18 +938,29 @@ def boresight_calibration(input_dataset, field_path='JWST_CALFIELD2020.csv', fie # call the target coordinates from the image header target_coordinate = (dataset[i].pri_hdr['RA'], dataset[i].pri_hdr['DEC']) + target_coord_tab = astropy.table.Table() + target_coord_tab['x'] = [np.shape(image)[1] // 2 ] # assume the target is at the center of the image + target_coord_tab['y'] = [np.shape(image)[0] // 2] + target_coord_tab['RA'] = [target_coordinate[0]] + target_coord_tab['DEC'] = [target_coordinate[1]] + target_coord_tables.append(target_coord_tab) + + # run automated source finder if field_matches are passed but distortion is also being computed - # run automated source finder if no matches are given - if field_matches is None: + if field_matches is not None and find_distortion is False: + matched_sources = matched_sources_multiframe[i] + else: found_sources = find_source_locations(image, threshold=find_threshold, fwhm=fwhm, mask_rad=mask_rad) matched_sources = match_sources(dataset[i], found_sources, field_path, comparison_threshold=comparison_threshold, rad=search_rad, platescale_guess=platescale_guess, platescale_tol=platescale_tol) + # matched_sources_multiframe.append(matched_sources) ## dont need to append these to the larger list because they are only used here for ps and na + # maybe there has to be a third case where we havent passed in matches but want distortion correction so we need to use the auto ones # compute the calibration properties cal_properties = compute_platescale_and_northangle(image, source_info=matched_sources, center_coord=target_coordinate, center_radius=center_radius) ra, dec = compute_boresight(image, source_info=matched_sources, target_coordinate=target_coordinate, cal_properties=cal_properties) # return a single AstrometricCalibration data file - astrom_data = np.array([ra, dec, cal_properties[0], cal_properties[1]]) + astrom_data = np.array([ra, dec, cal_properties[0], cal_properties[1], np.inf, np.inf]) astrom_cal = corgidrp.data.AstrometricCalibration(astrom_data, pri_hdr=dataset[i].pri_hdr, ext_hdr=dataset[i].ext_hdr, input_dataset=in_dataset) astroms.append(astrom_cal) @@ -735,13 +970,27 @@ def boresight_calibration(input_dataset, field_path='JWST_CALFIELD2020.csv', fie avg_platescale = np.mean([astro.platescale for astro in astroms]) avg_northangle = np.mean([astro.northangle for astro in astroms]) - avg_data = np.array([avg_ra, avg_dec, avg_platescale, avg_northangle]) + # compute the distortion map coeffs + if find_distortion: + first_stars, offsets, true_offsets, errs = format_distortion_inputs(input_dataset, matched_sources_multiframe, ref_star_pos=target_coord_tables, position_error=position_error) + distortion_coeffs, order = compute_distortion(input_dataset, first_stars, offsets, true_offsets, errs, platescale=avg_platescale, northangle=avg_northangle, fitorder=fitorder, initial_guess=initial_dist_guess) + else: + # set default coeffs to produce zero distortion + fitparams = (fitorder + 1)**2 + zero_dist = [0 for _ in range(fitorder+1)] + [500,] + [0 for _ in range(fitparams - fitorder - 2)] + [0,500] + [0 for _ in range(fitparams-2)] + distortion_coeffs = np.array(zero_dist) + order = fitorder + + astromcal_data = np.concatenate((np.array([avg_ra, avg_dec, avg_platescale, avg_northangle]), np.array(distortion_coeffs), np.array([order])), axis=0) + astroms_dataset = corgidrp.data.Dataset(astroms) - avg_cal = corgidrp.data.AstrometricCalibration(avg_data, pri_hdr=input_dataset[0].pri_hdr, ext_hdr=input_dataset[0].ext_hdr, input_dataset=astroms_dataset) + avg_cal = corgidrp.data.AstrometricCalibration(astromcal_data, pri_hdr=input_dataset[0].pri_hdr, ext_hdr=input_dataset[0].ext_hdr, input_dataset=astroms_dataset) # update the history history_msg = "Boresight calibration completed" astrom_cal_dataset = corgidrp.data.Dataset([avg_cal]) astrom_cal_dataset.update_after_processing_step(history_msg) + ## history message should be added to the input dataset(?) + input_dataset.update_after_processing_step(history_msg) - return astrom_cal \ No newline at end of file + return avg_cal \ No newline at end of file diff --git a/corgidrp/data.py b/corgidrp/data.py index 32eff9ef..42f53b6f 100644 --- a/corgidrp/data.py +++ b/corgidrp/data.py @@ -1311,7 +1311,9 @@ 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 a single array of calibration measurements of the following lengths (boresight: length 2 (RA, DEC), + plate scale: length 1 (float), north angle: length 1 (float), distortion coeffs: length dependent on order of polynomial fit but the last value should be an int describing the polynomial order). For a + 3rd order distortion fit the input array should be length 37. 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) @@ -1319,19 +1321,21 @@ class AstrometricCalibration(Image): 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), where the last value of the array is the order of polynomial used """ - 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:] # if this is a new astrometric calibration file, bookkeep it in the header # we need to check if it is new diff --git a/corgidrp/l2b_to_l3.py b/corgidrp/l2b_to_l3.py index 4f9650f2..b1370788 100644 --- a/corgidrp/l2b_to_l3.py +++ b/corgidrp/l2b_to_l3.py @@ -1,18 +1,69 @@ +import numpy as np +import astropy.wcs as wcs + # A file that holds the functions that transmogrify l2b data to l3 data -def create_wcs(input_dataset): +def create_wcs(input_dataset, astrom_calibration): """ Create the WCS headers for the dataset. Args: input_dataset (corgidrp.data.Dataset): a dataset of Images (L2b-level) + astrom_calibration (corgidrp.data.AstrometricCalibration): an astrometric calibration file for the input dataset Returns: corgidrp.data.Dataset: a version of the input dataset with the WCS headers added """ + updated_dataset = input_dataset.copy() - return input_dataset.copy() + northangle = astrom_calibration.northangle + platescale = astrom_calibration.platescale + center_coord = astrom_calibration.boresight + + # create wcs for each image in the dataset + for image in updated_dataset: + + im_data = image.data + image_shape = im_data.shape + center_pixel = [image_shape[1] // 2, image_shape[0] // 2] + roll_ang = image.pri_hdr['ROLL'] + + vert_ang = np.radians(northangle + roll_ang) ## might be -roll_ang + pc = np.array([[-np.cos(vert_ang), np.sin(vert_ang)], [np.sin(vert_ang), np.cos(vert_ang)]]) + cdmatrix = pc * (platescale * 0.001) / 3600. + + # create dictionary with wcs information + wcs_info = {} + wcs_info['CD1_1'] = cdmatrix[0,0] + wcs_info['CD1_2'] = cdmatrix[0,1] + wcs_info['CD2_1'] = cdmatrix[1,0] + wcs_info['CD2_2'] = cdmatrix[1,1] + + wcs_info['CRPIX1'] = center_pixel[0] + wcs_info['CRPIX2'] = center_pixel[1] + + wcs_info['CTYPE1'] = 'RA---TAN' + wcs_info['CTYPE2'] = 'DEC--TAN' + + wcs_info['CDELT1'] = (platescale * 0.001) / 3600 ## converting to degrees + wcs_info['CDELT2'] = (platescale * 0.001) / 3600 + + wcs_info['CRVAL1'] = center_coord[0] + wcs_info['CRVAL2'] = center_coord[1] + + wcs_info['PLTSCALE'] = platescale ## [mas] / pixel + + # update the image header with wcs information + for key, value in wcs_info.items(): + image.ext_hdr[key] = value + + # update the dataset with new header entries an dhistory message + history_msg = 'WCS created' + + updated_dataset.update_after_processing_step(history_msg) + + return updated_dataset def divide_by_exptime(input_dataset): """ diff --git a/corgidrp/l3_to_l4.py b/corgidrp/l3_to_l4.py index 3593d4b2..0956658f 100644 --- a/corgidrp/l3_to_l4.py +++ b/corgidrp/l3_to_l4.py @@ -1,26 +1,85 @@ # A file that holds the functions that transmogrify l3 data to l4 data from pyklip.klip import rotate +import scipy.ndimage from corgidrp import data from scipy.ndimage import rotate as rotate_scipy # to avoid duplicated name from scipy.ndimage import shift import numpy as np import glob -def distortion_correction(input_dataset, distortion_calibration): +def distortion_correction(input_dataset, astrom_calibration): """ Apply the distortion correction to the dataset. Args: input_dataset (corgidrp.data.Dataset): a dataset of Images (L3-level) - distortion_calibration (corgidrp.data.DistortionCalibration): a DistortionCalibration calibration file to model the distortion + astrom_calibration (corgidrp.data.AstrometricCalibration): an AstrometricCalibration calibration file to model the distortion Returns: corgidrp.data.Dataset: a version of the input dataset with the distortion correction applied """ + undistorted_dataset = input_dataset.copy() + distortion_coeffs = astrom_calibration.distortion_coeffs[0] + distortion_order = int(astrom_calibration.distortion_coeffs[1]) - return input_dataset.copy() + undistorted_ims = [] + + # apply the distortion correction to each image in the dataset + for undistorted_data in undistorted_dataset: + + im_data = undistorted_data.data + imgsizeX, imgsizeY = im_data.shape + + # set image size to the largest axis if not square imagea + if (imgsizeX >= imgsizeY): imgsize = imgsizeX + else: imgsize = imgsizeY + + yorig, xorig = np.indices(im_data.shape) + y0, x0 = imgsize//2, imgsize//2 + yorig -= y0 + xorig -= x0 + + ### compute the distortion map based on the calibration file passed in + fitparams = (distortion_order + 1)**2 + + # reshape the coefficient arrays + x_params = distortion_coeffs[:fitparams] + x_params = x_params.reshape(distortion_order+1, distortion_order+1) + + total_orders = np.arange(distortion_order+1)[:,None] + np.arange(distortion_order+1)[None,:] + x_params = x_params / 500**(total_orders) + + # evaluate the legendre polynomial at all pixel positions + x_corr = np.polynomial.legendre.legval2d(xorig.ravel(), yorig.ravel(), x_params) + x_corr = x_corr.reshape(xorig.shape) + + distmapX = x_corr - xorig + + # reshape and evaluate the same way for the y coordinates + y_params = distortion_coeffs[fitparams:] + y_params = y_params.reshape(distortion_order+1, distortion_order+1) + y_params = y_params /500**(total_orders) + + y_corr = np.polynomial.legendre.legval2d(xorig.ravel(), yorig.ravel(), y_params) + y_corr = y_corr.reshape(yorig.shape) + distmapY = y_corr - yorig + + # apply the distortion grid to the image indeces and map the image + gridx, gridy = np.meshgrid(np.arange(imgsize), np.arange(imgsize)) + gridx = gridx - distmapX + gridy = gridy - distmapY + + undistorted_image = scipy.ndimage.map_coordinates(im_data, [gridy, gridx]) + + undistorted_ims.append(undistorted_image) + + history_msg = 'Distortion correction completed' + + undistorted_dataset.update_after_processing_step(history_msg, new_all_data=np.array(undistorted_ims)) + + return undistorted_dataset def find_star(input_dataset): """ diff --git a/corgidrp/mocks.py b/corgidrp/mocks.py index cf278aae..2bd69ad3 100644 --- a/corgidrp/mocks.py +++ b/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,27 +987,92 @@ 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() prihdr['VISTYPE'] = 'BORESITE' prihdr['RA'] = target[0] prihdr['DEC'] = target[1] + prihdr['ROLL'] = 0 ## assume a telescope roll = 0 for now - newhdr = fits.Header(new_hdr) - frame = data.Image(sim_data, pri_hdr= prihdr, ext_hdr= newhdr) + # newhdr = fits.Header(new_hdr) + # frame = data.Image(sim_data, pri_hdr= prihdr, ext_hdr= newhdr) + ## save a default ext_hdr + frame = data.Image(sim_data, pri_hdr= prihdr, ext_hdr= exthdr) 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) diff --git a/tests/e2e_tests/astrom_e2e.py b/tests/e2e_tests/astrom_e2e.py index ff491ebd..31402ea1 100644 --- a/tests/e2e_tests/astrom_e2e.py +++ b/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") diff --git a/tests/test_astrom.py b/tests/test_astrom.py index c77179cb..d8a19d11 100644 --- a/tests/test_astrom.py +++ b/tests/test_astrom.py @@ -6,6 +6,7 @@ import corgidrp.mocks as mocks import corgidrp.astrom as astrom import corgidrp.data as data +import astropy.io.ascii as ascii def test_astrom(): """ @@ -31,7 +32,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 @@ -63,5 +63,122 @@ def test_astrom(): pickled_astrom = pickle.loads(pickled) assert np.all((astrom_cal.data == pickled_astrom.data)) # check it is the same as the original +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(astrom_cal.distortion_coeffs[-1]) + fitparams = (fitorder + 1)**2 + true_fitorder = int(expected_coeffs[-1]) + true_fitparams = (true_fitorder + 1)**2 + + # reshape the coeff arrays for the best fit and true coeff params + best_params_x = coeffs[: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[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 that the distortion error in the central 1" x 1" region (center ~45 x 45 pixels) + # has distortion error < 4 [mas] (~0.1835 [pixel]) + lower_lim, upper_lim = int((1024//2) - ((1000/21.8)//2)), int((1024//2) + ((1000/21.8)//2)) + central_1arcsec_x = x_diff[lower_lim: upper_lim+1,lower_lim: upper_lim+1] + central_1arcsec_y = y_diff[lower_lim: upper_lim+1,lower_lim: upper_lim+1] + true_1arcsec_x = true_x_diff[lower_lim: upper_lim+1,lower_lim: upper_lim+1] + true_1arcsec_y = true_y_diff[lower_lim: upper_lim+1,lower_lim: upper_lim+1] + + assert np.all(np.abs(central_1arcsec_x - true_1arcsec_x) < 0.1835) + assert np.all(np.abs(central_1arcsec_y - true_1arcsec_y) < 0.1835) + + # 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_astrom() \ No newline at end of file + test_astrom() + test_distortion() \ No newline at end of file diff --git a/tests/test_create_wcs.py b/tests/test_create_wcs.py new file mode 100644 index 00000000..88cf5c9d --- /dev/null +++ b/tests/test_create_wcs.py @@ -0,0 +1,65 @@ +import numpy as np +import os +from corgidrp import mocks, astrom +from corgidrp.l2b_to_l3 import create_wcs + +def test_create_wcs(): + """ + Unit test of the create WCS function. The test checks that all WCS keywords (['CD{}_{}'], ['CRPIX{}'], ['CTYPE{}'], ['CDELT{}'], ['CRVAL{}']) exist in the updated ext_hdr + and that the values are consistent with those derived from the AstrometricCalibration file. + """ + + # create mock dataset (arbitrary northangle) + field_path = os.path.join(os.path.dirname(__file__), "test_data", "JWST_CALFIELD2020.csv") + mock_dataset = mocks.create_astrom_data(field_path, platescale=21.8, rotation=20) + + # run the boresight calibration to get an AstrometricCalibration file + astrom_cal = astrom.boresight_calibration(mock_dataset, field_path, find_threshold=100) + + # create the wcs + updated_dataset = create_wcs(mock_dataset, astrom_cal) + + # check that all wcs keywords exist in the ext_hdrs of the mock dataset + # and that the values are as expected from the AstrometricCalibration file + platescale = astrom_cal.platescale + northangle = astrom_cal.northangle + boresight = astrom_cal.boresight + + for mock_frame, updated_frame in zip(mock_dataset, updated_dataset): + roll_ang = mock_frame.pri_hdr['ROLL'] + data = mock_frame.data + image_shape = data.shape + center_pixel = [image_shape[1] // 2, image_shape[0] // 2] + + pc = np.array([[-np.cos(np.radians(northangle + roll_ang)), np.sin(np.radians(northangle + roll_ang))], [np.sin(np.radians(northangle + roll_ang)), np.cos(np.radians(northangle + roll_ang))]]) + matrix = pc * (platescale * 0.001) / 3600. + + # gather expected values in a dictionary + expected = {} + expected['CD1_1'] = matrix[0,0] + expected['CD1_2'] = matrix[0,1] + expected['CD2_1'] = matrix[1,0] + expected['CD2_2'] = matrix[1,1] + + expected['CRPIX1'] = center_pixel[0] + expected['CRPIX2'] = center_pixel[1] + + expected['CTYPE1'] = 'RA---TAN' + expected['CTYPE2'] = 'DEC--TAN' + + expected['CDELT1'] = (platescale * 0.001) / 3600 ## converting to degrees + expected['CDELT2'] = (platescale * 0.001) / 3600 + + expected['CRVAL1'] = boresight[0] + expected['CRVAL2'] = boresight[1] + + # gather the wcs values from the updated dataset + wcs = {} + for key in expected.keys(): + wcs[key] = updated_frame.ext_hdr[key] + + # compare the expected dictionary to the updated dateset output + assert wcs.items() == expected.items() + +if __name__ == "__main__": + test_create_wcs() \ No newline at end of file diff --git a/tests/test_data/distortion_expected_coeffs.csv b/tests/test_data/distortion_expected_coeffs.csv new file mode 100644 index 00000000..7b26b360 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/test_distortion_correction.py b/tests/test_distortion_correction.py new file mode 100644 index 00000000..6e011437 --- /dev/null +++ b/tests/test_distortion_correction.py @@ -0,0 +1,35 @@ +import numpy as np +import pytest +import os + +from corgidrp import data, mocks, astrom +from corgidrp.l3_to_l4 import distortion_correction + +def test_distortion_correction(): + """ + Unit test of the distortion correction function. The test checks that a distorted mock image has been accurately corrected for by comparing it to an undistorted mock image. + + """ + + # create an unditorted dataset + field_path = os.path.join(os.path.dirname(__file__), "test_data", "JWST_CALFIELD2020.csv") + no_distortion_dataset = mocks.create_astrom_data(field_path) + + # add distortion to the dataset + distortion_coeffs_path = os.path.join(os.path.dirname(__file__), "test_data", "distortion_expected_coeffs.csv") + distortion_coeffs = np.genfromtxt(distortion_coeffs_path) + distortion_dataset = mocks.create_astrom_data(field_path, distortion_coeffs_path=distortion_coeffs_path) + + # assume a ground truth AstrometricCalibration file + astromcal_data = np.concatenate((np.array([80.553428801, -69.514096821, 21.8, 45]), distortion_coeffs), axis=0) + astrom_cal = data.AstrometricCalibration(astromcal_data, pri_hdr=distortion_dataset[0].pri_hdr, ext_hdr=distortion_dataset[0].ext_hdr, input_dataset=distortion_dataset) + + # use the distortion correction function to undistort + undistorted_dataset = distortion_correction(distortion_dataset, astrom_cal) + + # compare the undistorted data to the original dataset with no distortion + for frame, ref_frame in zip(undistorted_dataset, no_distortion_dataset): + assert np.mean(frame.data - ref_frame.data) == pytest.approx(0, abs=0.005) + +if __name__ == "__main__": + test_distortion_correction() \ No newline at end of file