ovro-eovsa / ovro-lwa-solar

Calibration, imaging, and analysis of solar data taken by the Owens Valley Long Wavelength Array (OVRO-LWA)
MIT License
5 stars 6 forks source link

Developing refraction correction techniques #48

Open dgary50 opened 9 months ago

dgary50 commented 9 months ago

Here is the code (updated 2023-12-11) to correct an fch file and write out a new fits file result. This version corrects a critical error from before, and also creates and writes out new mfs files. Note that the refraction 1/f^2 fit is actually done using the mfs files, so both are needed. To use it, create a new folder to store the results, cd to that folder, and then use the command: import mfits2mfs as m refraction = m.mfits2mfs(filelist) where filelist is the full list of files (expected to be all of those for a given date, although any list of files will do). The list should be of the fch files, but the mfs files of the same name are also needed. If you give it a single filename, it still needs to be a list, i.e. [filename]. The refraction dictionary returned has keys 'time', 'refraction', 'offset'. The later two are n x 2 arrays that are basically the fit parameters for x and y returned by the linear fit to 1/f^2, where f is in MHz. Therefore, the 'refraction' values seem high, but to get them for, say, 50 MHz, just divide by 50^2. Then they will be the refraction in arcmin for 50 MHz.

It is helpful to visualize the result for a given file by m.check_as_movie(filename). This is very slow from NJ, so a quicker check is with m.check_as_contour(filename). With luck, when called with the name of one of the corrected files, the image will be stationary with frequency (although not necessarily centered).

import numpy as np
import matplotlib.pylab as plt
from matplotlib import colormaps
import os
from glob import glob
from skimage.registration import phase_cross_correlation
from astropy.io import fits

def threshold_image(img, per_thrshd=98):
    """
    Set pixels with intensity below 95 percentile to 0.

    Parameters:
    - img: numpy ndarray, the input image

    Returns:
    - thresholded_img: numpy ndarray, the thresholded image
    """
    percentile_img = np.percentile(img, per_thrshd)
    thresholded_img = np.where(img < percentile_img, 0, img)
    return thresholded_img

def rd_mfits(fname,reverse=True):
    ''' Given a multifrequency file name, return the frequencies, images, and header.
        If reverse is True (default), the order of images and frequencies are reversed
        to range from high to low (since it is more natural to start at high frequencies
        for determining shifts.
    '''
    imgs = fits.getdata(fname)
    h = fits.getheader(fname)
    f_recs = fits.getdata(fname, 1)
    freqs = [f[0] / 1e6 for f in f_recs]
    freqs = np.array(freqs)
    imgs = imgs.squeeze()
    if reverse:
        #print('Frequency order will be high to low')
        freqs = freqs[::-1]
        imgs = imgs[::-1]
    return freqs, imgs, h

def shifts_freq_DG(imgs):
    ''' Given a multifrequency imagecube, determine the accumulated pixel shifts
        needed to align the images, aligning each image(i) to image(i-1). 
        Return the result as a 2D array, the second index being 0 for x and 1 for y.
    '''
    n = len(imgs)
    accumulated_shift = [np.zeros(2)] * n
    for i in range(1,n):
        shift, _, _ = phase_cross_correlation(imgs[i-1], imgs[i], upsample_factor=10)
        accumulated_shift[i] = accumulated_shift[i-1] + shift
    shiftpx = np.array(accumulated_shift)
    return shiftpx

def shifts_freq_SJ(imgs, idx_gap=2):
    ''' Given a multifrequency imagecube, determine the accumulated pixel shifts
        needed to align the images, aligning each image(i) to min(i+n/2,n-1). In
        other words, if n = 120, image(0) is aligned with image(60), 
        image(1) with image(61), and so-on until image(60), then align all the
        rest with the highest frequency image...

        Return the result as a 2D array, the first index being 0 for x and 1 for y.

        My feeling is that this might have worked well for the eclipse, where n = 12,
        but for large n it might be better to start with the highest frequency
        and work backwards, just aligning each image with its immediate neighbor.
    '''
    n = len(imgs)
    accumulated_shift = [np.zeros(2)] * (n - 1)

    #idx_gap = n // 2
    for i, img in enumerate(imgs[1:]):
        ref_index = max((i - idx_gap, 0))
        ref_image = imgs[ref_index]
        shift, _, _ = phase_cross_correlation(ref_image, img, upsample_factor=10)
        if ref_index == n - 1:  # if reference image is the last image
            accumulated_shift[i] = shift
        else:  # if reference image is not the last image
            accumulated_shift[i] = shift + accumulated_shift[i - idx_gap]

    shiftpx = np.array(accumulated_shift)
    return shiftpx

def shifts_tim_SJ(imgs, ref_image):
    ''' Given a multitime imagecube at a given frequency, determine the accumulated
        pixel shifts, aligning each image with its neighbor "idx_gap" away.

        Return the result as a 2D array, the second index being 0 for x and 1 for y.
    '''
    n = len(imgs)
    accumulated_shift = [np.zeros(2)] * n

    for i, img in enumerate(imgs):
        #ref_index = min(i + idx_gap, n - 1)
        # print(f'i: {i} ---  ref: {ref_index}')
        #ref_image = imgs[ref_index]
        shift, _, _ = phase_cross_correlation(ref_image, img, upsample_factor=20)
        # if ref_index == n - 1:  # if reference image is not the last image
        #     accumulated_shift[i] = shift
        # else:  # if reference image is the last image
        accumulated_shift[i] = shift # + accumulated_shift[i - idx_gap]
    #accumulated_shift[-1] = accumulated_shift[-2] + shift
    shiftpx = np.array(accumulated_shift)
    #shiftpx = np.roll(shiftpx, 1, axis=0)
    #shiftpx[0] = 0
    return shiftpx

def register(freqs, shiftpx, nskip=40, order=1, showplt=False):
    ''' Given a multifrequency datacube of images, and pixel shifts from 
        shifts_freq_SJ(), determine the best linear fit vs. f**(-2) and
        return the fitting parameters px, py.  They are just slope and 
        offset for each axis.  Optionally plot the results of the fits.

        Pass in the filename in showplt to use its time for a plot title.
    '''
    if2 = freqs ** (-2)
    px = np.polyfit(if2[nskip:], shiftpx[nskip:, 0], order)
    py = np.polyfit(if2[nskip:], shiftpx[nskip:, 1], order)
    if showplt:
        f, ax = plt.subplots(1, 1)
        ax.plot(if2, shiftpx[:, 0], '.', label='xshift')
        ax.plot(if2, shiftpx[:, 1], '.', label='yshift')
        fq = np.linspace(0, 0.001, 200)
        ax.plot(fq, np.polyval(px, fq), color='C0')
        ax.plot(fq, np.polyval(py, fq), color='C1')
        ax.text(-0.0001, px[-1], str(px[-1])[:4], color='C0')
        ax.text(-0.0001, py[-1], str(py[-1])[:4], color='C1')
        ax.set_xlim(-0.0002, 0.0010)
        ax.set_ylim(-50, 50)
        if type(showplt) is str: ax.set_title('Time:'+showplt.split('T')[1].split('Z')[0])
        plt.legend()
    return px, py

def apply_shift(freqs, imgs, px, py):
    ''' Given x and y pixel shifts vs. frequency output by register(), actually apply 
        the shifts and return the shifted images.    
    '''
    if2 = freqs ** (-2)
    xshift = np.polyval(px, if2)
    yshift = np.polyval(py, if2)
    imgsout = np.zeros_like(imgs)
    imgsout[0] = imgs[0]
    for i, img in enumerate(imgs[1:]):
        imgsout[i+1] = np.roll(np.roll(img, round(xshift[i]), axis=0), round(yshift[i]), axis=1)
    return imgsout

def fits_corr(fitsfile):
    ''' Given a multifrequency fits filename, read the images and determine the shifts
        for a subset of frequencies (starting with the 40th frequency and skipping every 5th
        frequency).  Then determine the best fit line to the shifts vs. f**(-2), and finally
        apply the shifts, returning the corrected images.
    '''
    freqs, imgs = rd_mfits(fitsfile)  # Reads the multifits file
    slc = slice(40, -1, 5)
    img2proc = imgs[slc]

    img2proc = [threshold_image(img, per_thrshd=97) for img in img2proc]

    shifts = shifts_freq_SJ(img2proc)

    px, py = register(freqs[slc], imgs[slc], shifts, nskip=6, order=1)  # Determines fit to shifts
    imgsout = apply_shift(freqs, imgs, px, py)  # Applies shifts and returns shifted images

    return imgsout, px, py

def mfits2mfs(files,showplt=False):
    ''' Reads a list of fch (frequency-channel) files and does a series of
        steps for each file:
           1. Reads the mfs images corresponding to the file. (rd_mfits)
           2. Determines the shifts vs. frequency using phase cross correlation. (shifts_freq_SJ)
           3. Fits a linear 1/f**2 dependence to the higher frequencies (those > 45 MHz) (register)
           4. Reads the fch (many-frequency) file and applies the fit parameters (apply_shift)
           5. Writes the corrected fch images to a new file in the current directory.
        The following are yet to do:
           6. Combine the best images of each band frequency range to create a pseudo-mfs image.
           7. Read the corresponding msf fits file, replace the mfs images and save them.
           8. Create new 12-panel png file with the corrected mfs images.
    '''
    from time import time
    # Read in mfs fits files:
    refraction = []
    offset = []
    times = []
    for i,file in enumerate(files):
        print('Working on file',i,'of',len(files),end='')
        t1 = time()
        # Note, rd_mfits() defaults to reversing the frequency order (high to low)
        mfreqs, mimgs, h = rd_mfits(file.replace('fch','mfs'))
        times.append(h['DATE-OBS'])
        img2proc = [threshold_image(img, per_thrshd=95) for img in mimgs]
        shiftpx = shifts_freq_DG(img2proc)
        good, = np.where(mfreqs > 45.0)
        px, py = register(mfreqs[good], shiftpx[good], nskip=0, showplt=showplt)
        refraction.append([px[0],py[0]])
        offset.append([px[1],py[1]])
        freqs, imgs, h = rd_mfits(file)    # Note, images are in reverse freq. order
        imgsout = apply_shift(freqs, imgs, px, py)
        # Form new mfs images from the mean of the shifted images
        for i in range(len(mfreqs)):
            idx, = np.where(np.logical_and(freqs < (mfreqs[i]+2.297), freqs >= (mfreqs[i]-2.297)))
            mimgs[i] = np.nanmean(imgsout[idx],0)
        # Write out new fch fits images
        hdu1 = fits.open(file)
        fname = os.path.basename(file)
        imgs = hdu1[0].data
        imgs[0] = imgsout[::-1]    # Unreverse the order when writing back out.
        hdu1.writeto(fname, overwrite=True)
        hdu1.close()
        # Write out new mfs fits images
        hdu1 = fits.open(file.replace('fch','mfs'))
        fname = os.path.basename(file.replace('fch','mfs'))
        imgs = hdu1[0].data
        imgs[0] = mimgs[::-1]    # Unreverse the order when writing back out.
        hdu1.writeto(fname, overwrite=True)
        hdu1.close()
        print(' ...',str(time()-t1)[:5],'s')

    refr = {'time':np.array(times), 'refraction':np.array(refraction), 'offset':np.array(offset)}
    # Save the output in the same folder as the data files
    np.save('refraction_'+refr['time'][0][:10].replace('-','')+'_save.npy', refr)
    return refr

def check_as_movie(filename):
    # Given a filename (either corrected or uncorrected), play a simple movie to see
    # the multifrequency images.
    freqs, imgs, h = rd_mfits(filename)
    plt.clf()
    im = plt.imshow(imgs[0])
    for img in imgs[1:]:
        im.set_data(img)
        plt.pause(0.1)

def check_as_contour(filename,nmax=None):
    # Given a filename (either corrected or uncorrected), play a simple movie to see
    # the multifrequency images.
    freqs, imgs, h = rd_mfits(filename)
    if nmax is None:
        nmax = len(freqs)
    plt.clf()
    im = plt.imshow(imgs[0],origin='lower')
    cmap = colormaps['Spectral']
    for i,img in enumerate(imgs[1:nmax]):
        plt.contour(img,[np.max(img)/2.],colors=[cmap(i/nmax)])

def refr_plot(datstr):
    ''' Plots the magnitude of refraction for a given date, as provided by the
        corresponding refraction save (npy) file, in a standard format (for
        comparison from day to day).  This writes the plot as a png file in
        the same folder selected by datstr.

        Inputs:
          datstr   A text string as YYYYMMDD.
    '''
    from matplotlib.dates import DateFormatter
    from astropy.time import Time
    os.chdir('/data07/dgary/refraction_tests/'+datstr)
    blah = np.load('refraction_'+datstr+'_save.npy',allow_pickle=True)
    refr = blah.item()
    pd = Time(refr['time']).plot_date
    fig, ax = plt.subplots(1,1,figsize=(12,3))
    ax.plot_date(pd,np.sqrt((refr['refraction'][:,0]+200*85**2)**2 + refr['refraction'][:,1]**2)/85**2-200,'.')
    ax.set_title(datstr[:4]+'-'+datstr[4:6]+'-'+datstr[6:]+' Refraction Shift Magnitude [arcmin at 85 MHz]')
    ax.set_xlabel('Time [UTC]')
    ax.set_ylabel('Shift Magnitude [arcmin]')
    ax.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
    ax.set_ylim(-80,80)
    ax.set_xlim(int(pd[0])+0.54167,int(pd[0])+1.083333)
    plt.savefig('refraction_'+datstr+'.png')

def show_all(filename):
    ''' Shows all images in a file on a single plot.
    '''
    import matplotlib.ticker as ticker
    freqs, imgs, h = rd_mfits(filename)
    nf = len(freqs)
    if nf <= 12:
        f, ax = plt.subplots(3,4,figsize=(5.25,4.6))
        plt.subplots_adjust(top=0.88, bottom=0.11, left=0.05, right=0.95, hspace=0.0, wspace=0.0)
    else:
        f, ax = plt.subplots(12,12,figsize=(10,10))
        plt.subplots_adjust(top=0.981, bottom=0.038, left=0.036, right=0.983, hspace=0.0, wspace=0.0)
    ax = ax.flatten()
    for i in range(nf):
        ax[i].imshow(imgs[i,32:256-32,32:256-32],origin='lower')
        ax[i].xaxis.set_major_locator(ticker.NullLocator())
        ax[i].yaxis.set_major_locator(ticker.NullLocator())
        ax[i].text(5,165,str(freqs[i])[:5],color='w',fontsize=8)
dgary50 commented 9 months ago

Here is the code to create the ionospheric gravity wave model. The density Ne is returned. Here is my presentation, with more details of the model: Ionospheric_Refraction.pptx

dtor = np.pi/180.

def refraction_model(t=0, T=40, gamma=1):
    ''' Based on Koval et al. (2018), Journal of Geophysical Research: Space Physics, 
        123, 8940–8950. https://doi.org/10.1029/2018JA025584.  

        Creates a 1001 x 1001 pixel density model of the disturbed
        ionosphere, with vertical resolution 0.5 km, horizontal resolution 1 km.

        Inputs:
           t       The time (seconds) at which to calculate the model.  
                     Calling the routine with changing time t allows the waves 
                     in the model to propagate.
           T       The period of the gravity wave disturbance (minutes)
           gamma   A knob to adjust the magnitude of the disturbance.

        Many other parameters are hard-coded in the first section below, but could be 
        converted to setable parameters on calling the routine, if desired.
    '''
    # Chosen parameter values for the problem
    lambda_x = 300  # [km] horizontal wavelength of gravity wave
    A = 135*dtor    # Azimuth of gravity wave travel
    dip = 61.5*dtor # Dip angle of local magnetic field
    u_parallel = 7e-3  # [km/s] component of neutral wind velocity along geomagnetic field
    U = 0           # Background wind velocity
    N_max = 1.24e12 # [m^-3] Maximum density of F-layer
    z_max = 300     # [km] Height of peak F-layer density
    H = 50          # [km] Scale height for F-layer
    T_BV = 12       # Brunt-Vaisala period of buoyancy oscillation
    B_dec = 0 ovro-eovsa/ovro-lwa-solar#12*dtor # Magnetic declination (direction to magnetic pole)

    # Calculated quantities
    omega = 2*np.pi/(T*60.)         # Angular frequency of gravity wave (from period T)
    kx  = 2*np.pi/lambda_x          # horizontal component of gravity wave number (from lambda_x)
    c = omega/kx                    # Phase velocity of gravity wave (should be omega/k?)
    omega_B = 2*np.pi/(T_BV*60)     # Brunt-Vaisala frequency (from T_BV)
    z = np.linspace(0, 500, 1001)   # Height range of ionosphere (resolution is 1/2 km)
    x = np.linspace(0, 1000, 1001)  # Horizontal range (resolution is 1 km)
    kz2 = (omega_B/(U-c))**2 - kx**2 - 1/(4*H**2)
    kz = np.sqrt(kz2)               # Vertical wavenumber of gravity wave
    # k_parallel is the component of the wave number of the gravity wave along B.  It has a vertical
    # component k_parallelv, and a horizontal component k_parallelh.
    slope = np.arctan(kz/kx)
    print('Altitude of trough is:', 90 - slope/dtor, 'degrees')
    k_parallelv = np.sqrt(kx**2 + kz2) * np.cos(slope - dip)*np.cos((np.pi - A) + B_dec)
    k_parallelh = np.sqrt(kx**2 + kz2) * np.cos(slope - dip)*np.cos((np.pi - A) - B_dec)

    # These are all a function of z (height)    
    Ne0 = N_max*np.exp(0.5*(1-(z-z_max)/H - np.exp(-(z-z_max)/H)))      # Chapman density function
    dNe0dz = Ne0/(2*H) * (np.exp(-(z-z_max)/H) - 1)                     # Derivative of Ne0
    Nep = 1j*u_parallel/omega*(dNe0dz*np.sin(dip) - 1j*k_parallelv*Ne0) # Density perturbation

    # Calculate 2D density
    Ne = np.zeros((1001,1001), np.float64)
    for i,xi in enumerate(x):
        Ne[:,i] = Ne0 + gamma*np.abs(Nep)*np.cos(omega*t - kx*xi - kz*z)
    return Ne
surajit4444mondal commented 9 months ago

I have modified the rd_mfits function a bit to make it more general such that it can read the compressed fitsfiles produced by ndfits as well and am attaching it here.

def rd_mfits(fname,reverse=True):
    ''' Given a multifrequency file name, return the frequencies, images, and header.
        If reverse is True (default), the order of images and frequencies are reversed
        to range from high to low (since it is more natural to start at high frequencies
        for determining shifts.
    '''

    hdu=fits.open(fname)

    try:
      num_hdu=len(hdu)
      if num_hdu==2:
        imgs = hdu[0].data
        h = hdu[0].header
        f_recs = hdu[1].data
      elif num_hdu==3 and hdu[1].name=='COMPRESSED_IMAGE':
        imgs = hdu[1].data
        h = hdu[1].header
        f_recs = hdu[2].data
    finally:
      hdu.close()
    freqs = [f[0] / 1e6 for f in f_recs]
    freqs = np.array(freqs)
    imgs = imgs.squeeze()
    if reverse:
        #print('Frequency order will be high to low')
        freqs = freqs[::-1]
        imgs = imgs[::-1]
    return freqs, imgs, h
binchensun commented 6 months ago

For our record, @surajit4444mondal @peijin94 could you both post your updates on the methods you use here?

peijin94 commented 6 months ago

(from the dev notes)

Center alignment from low Tb contour

refraction correction

Thresh

The thresh needs to be low enough to include all the corona weak emission, and be less influenced by the on disk emission.

So we are using fixed value of 0.1 T_b(quiet Sun), empirical expression from Zhang et al 2022

def thresh_func(freq): # freq in Hz
    return 1.1e6 * (1-1.8e4* freq**(-0.6)) 

thresh = thresh_func*0.1

Remove bad artifacts

Using morphological operation to achieve a relatively robust way to get the coronal region while reducing the influence of artifacts to the maximum extend.

Untitled

The resulting correction is OK-ish

But in the correction displacement, there is a feature.

In the freq-displacement plot, we see the overall frequency-dependent trend looks like refraction, yes

Untitled 1

But the three jump near 43MHz 61MHz and 80MHz, it is obviously not refraction.

Before and after

Before

Untitled 2

After

Untitled 3

The source position jump is visible by eye:

Untitled 4

podman pull https://hub.docker.com/r/astronrd/linc

More test:

Quiet Sun with an active region

Untitled 5

when it is very faint:

Untitled 6

When it is too bright…..

Untitled 7

offset with time (x and y):

Untitled 8 Untitled 9

blue: x[46MHz]-x[76Mhz] , orange: y[46MHz]-y[76Mhz]

without and with fitting: Untitled 10 Untitled 11

fitting to x = x0+ a (1/f^2) and y = y0+ b(1/f^2)

x0 and y0: Untitled 12

Indeed there is frequency independent shift… (x0 y0)

binchensun commented 3 months ago

Refraction correction implemented in PRs #50, #54, #65, #67, #69, #72. Also implemented in the realtime pipeline with interpolation in ovro-lwa-solar-ops PRs https://github.com/ovro-eovsa/ovro-lwa-solar-ops/pull/16 and https://github.com/ovro-eovsa/ovro-lwa-solar-ops/pull/17.

binchensun commented 3 months ago

New feature request: divide the refraction into a constant and a variable component and at least correct the constant one determined from the whole day.

dgary50 commented 3 months ago

Depending on the number of good shift measurements and their temporal behavior, fitting linear or quadratic terms may be appropriate, too.

binchensun commented 3 months ago

Depending on the number of good shift measurements and there temporal behavior, fitting linear or quadratic terms may be appropriate, too.

This has already been implemented as a keyword (using scipy's interp1d function). See https://github.com/ovro-eovsa/ovro-lwa-solar/blob/main/ovrolwasolar/refraction_correction.py#L277. linear interpolation is used as the default. However, how to decide on the exact interpolation method remains an open question (which has not been thought out yet).