EUI to AIA remapping near sun-earth line

eui_to_aia.py

import numpy as np 
import sunpy
import imreg_dft

import astropy
from astropy import units as u
from astropy.wcs import WCS
from reproject import reproject_adaptive

def remap_eui(eui_map, aia_map, aia_file, out_shape = (256,256)):
    """remap eui onto aia, and degrade eui to aia resolution

    Parameters
    ----------
    eui_map : sunpy.map.Map
        HRT sunpy map object.
    aia_map : sunpy.map.Map
        HMI sunpy map object.
    out_shape : tuple
        Optional output size of remapped EUI

    Returns
    -------
    aia_overlap : sunpy.map.Map
        AIA sunpy input map. 
    eui_remap : sunpy.map.Map
        EUI remapped sunpy map object.
    """
    # define new header for eui map using aia observer coordinates
    out_header = sunpy.map.make_fitswcs_header(
        out_shape,
        aia_map.reference_coordinate.replicate(rsun=eui_map.reference_coordinate.rsun),
        scale=u.Quantity(aia_map.scale),
        instrument="EUI",
        observatory="SOLO",
        wavelength=eui_map.wavelength
    )

    #set remaining WCS keywords
    out_header['dsun_obs'] = aia_map.coordinate_frame.observer.radius.to(u.m).value
    out_header['hglt_obs'] = aia_map.coordinate_frame.observer.lat.value
    out_header['hgln_obs'] = aia_map.coordinate_frame.observer.lon.value

    aia_hdr = fits.open(aia_file)[1].header
    out_header['CROTA'] = aia_hdr['crota2']

    out_wcs = WCS(out_header)
    
    # reprojection
    output, footprint = reproject_adaptive(eui_map, out_wcs, out_shape)
    eui_remap = sunpy.map.Map(output, out_header)

    return (aia_map,eui_remap)

def eui_to_aia_and_align(eui_map, aia_map, aia_file):
    """remap eui onto aia, degrading eui to aia, and performing one iteration of shifting the AIA post EUI reproject to better align

    Parameters
    ----------
    eui_map : sunpy.map.Map
        HRT sunpy map object.
    aia_map : sunpy.map.Map
        HMI sunpy map object.
    aia_file: str
        path to AIA file (needed to open secondary header)

    Returns
    -------
    aia_overlap : sunpy.map.Map
        AIA sunpy map with direct overlap with EUI remapped. 
    eui_logpol : sunpy.map.Map
        EUI remapped sunpy map object with one polar log transform alignment procedure.
    """
    #take the submap of the aia using eui coords
    top_right = aia_map.world_to_pixel(eui_map.top_right_coord)
    bottom_left = ara_map.world_to_pixel(eui_map.bottom_left_coord)

    tr = np.array([top_right.x.value,top_right.y.value])
    bl = np.array([bottom_left.x.value,bottom_left.y.value])
    aia_og_size =  aia_map.submap(bl*u.pix,top_right=tr*u.pix)
    
    #add padding to aia submap so that extra pixels left when shifting later to better match to eui
    pad_tr = tr + 200
    pad_bl = bl - 200
    aia_pad = aia_map.submap(pad_bl*u.pix,top_right=pad_tr*u.pix)
    
    #remap eui with WCS
    aia_shape = aia_og_size.data.shape
    _, eui_remap = remap_eui_og_scale(eui_map, aia_og_size, aia_file, out_shape = aia_shape)
 
    #get residual rotation/shift using EUI as the reference
    s = int(eui_remap.data.shape[0]*0.2) # ignore boundaries when shifting
    imref = eui_remap.data[s:-s,s:-s]
    imtemp = aia_og_size.data[s:-s,s:-s]
    r = imreg_dft.similarity(imref,imtemp, 
                             numiter=3,constraints=dict(scale=(1,0)))
   
    #polar log transform on aia image to better align to eui
    aia_shifted = imreg_dft.transform_img(aia_pad.data,scale=1,angle=r['angle'],tvec=r['tvec'])
    
    h = aia_pad.fits_header
    h.append(('SHIFTX',r['tvec'][1],'shift along X axis (AIA-pixel)'),end=True)
    h.append(('SHIFTY',r['tvec'][0],'shift along Y axis (AIA-pixel)'),end=True)
    h.append(('RANGLE',r['angle'],'rotation angle (deg)'),end=True)
    print(r['angle'], r['tvec'][1], r['tvec'][0])
    aia_logpol = sunpy.map.Map((aia_shifted,h))
    
    #submap on the padded aia image to get the real overlap aia image back
    top_right = aia_logpol.world_to_pixel(eui_remap.top_right_coord)
    bottom_left = aia_logpol.world_to_pixel(eui_remap.bottom_left_coord)
    tr = np.array([top_right.x.value,top_right.y.value])
    bl = np.array([bottom_left.x.value,bottom_left.y.value])
    aia_overlap = aia_logpol.submap(bl*u.pix,top_right=tr*u.pix)
    
    return aia_overlap, eui_remap