AlphaScreen.py

'''
Obtained from EFY, SWRI
Containing functions that would be used for constructing the alpha screen.
'''

import numpy as np
from scipy.integrate import cumtrapz
from scipy.integrate import trapz
from scipy.interpolate import interp1d
from scipy.ndimage.interpolation import zoom


# 29-JUN-2019 (EFY, SWRI)
# Routines in this file generate an aperture (phase screen) from an atmospheric T(z) profile
# and a surface pressure. Assumes a N2 atmosphere.
# THIS APERTURE DIFFERS from the ones in the PLHaze_v40.py file: it uses Goodman's
# method of building the aperture in the Fresnel regime
# - Pick the FOV and the number of points
# - Build the aperture at the resolution specified (npts/FOV)
# In the PLHaze file, we used the Tretoer algorithm, which determined the FOV from
# the wavelength, distance to the aperture and the number of points, as follows:
# FOV = sqrt(npts * lam* D)
# Obviously, the Goodman aperture is more flexible.

# --------------------------------------------

def T2nu(r, Tr, gSurf, rSurf, PSurf, lam):
    '''Calculates the refractivity as a function of r from the T(r) temperature profile.
    Assumes hydrostatic equilibrium and the ideal gas law.

    r = radius (km)
    Tr = temperatures at r's (K)

    gSurf = acc. due to gravity at the surface (cgs; gSurf = 65.8 cm/sec^2 on Pluto)
    rSurf = radius of Pluto's surface (km)
    PSurf = surface pressure (cgs = microbars; 10.613 microbars for Pluto)

    m = molecular weight (g/mol); m = 28 for N2 gas. WE ASSUME N2!!!

    lam = wavelength (microns)

    EFY, SwRI, 21-MAY-2013
    '''

    k_Bolt = 1.3807e-16  # cgs, Boltzmann's constant
    Avo = 6.022e23  # molecules per mole (Avogadro's number)

    m = 28.0

    r_cm = r * 1.0e5  # radii in cm

    gr = gSurf * (rSurf / r) ** 2  # gr is the acceleration due to grav. at r.
    Hr = (k_Bolt * Tr) / ((m / Avo) * gr)  # scale height as a function of r.

    # Now integrate to get pressure, P(r). We'll use the relation:
    # ln(P) = integral(-1/H)dr

    lnP1 = cumtrapz(-1.0 / Hr, r_cm)
    lnP = np.concatenate([np.zeros(1), lnP1])  # We want a leading zero at P(rSurf)

    Pn = np.exp(lnP)
    # Pn is normalized. We need to scale it by a factor of PSurf/Pn(rSurf)
    Pf = interp1d(r, Pn, kind='linear', fill_value='extrapolate')
    Pr = Pn * (PSurf / Pf(rSurf)) # / 1.0e6 # unit conversion from microbar to bar

    # Now get the refractivity as a function of height

    nu_r = nu_ref_N2(Pr, Tr, lam)

    return (Pr, nu_r)


# --------------------------------------------

def TP2nu(r, Tr, Pr, lam):
    '''
    Like T2nu(), expect that this function TAKES P(r) as an ARGUMENT instead
    of calculating it from the T(r) profile. The refractivity vs. r is returned.

    r = radius (km
    Tr = temperatures at r values (K)
    Pr = pressures at r values (cgs: microbars)

    lam = wavelength (microns)

    EFY, SwRI, 2-MAY-2015

    '''

    nu_r = nu_ref_N2(Pr, Tr, lam)

    return (nu_r)


# --------------------------------------------

def nu2alpha(r, rSurf, nu_r, rMax, dr, sMax, ds):
    '''This routine takes the refractivity profile (nu vs. r) as input and returns the
    integrated line-of-sight refractivity (alpha vs. r) as output.

    Definition: Impact parameter => for a chord through the atmosphere, the impact parameter is
    the radius at which the chord passes closest to the solid surface.

    INPUTS:
    r = radii (km)
    rSurf = surface radius (km)
    nu_r = refractivities vs. r
    rMax = the highest radius (km, aka the impact parameter) for which we'll calculate alpha.
    dr = the length of step size for sampling the impact parameters (km).
    sMax = the length of the line of sight chord from the midpoint (km)
    ds = the length of an integration step (km) along a chord through the atmosphere.

    OUTPUTS:
    r_impVec: radii (km) of impact parameters at which we calculate alpha0
    alpha0: the integrated line-of-sight refractivities evaluated at each r_impact.
    aFun: the function that returns alpha as a function of an impact parameter (r)

    EFY, SWRI, 29-JUN-2019    '''

    # STEP 1: build a function to give us nu as a function of r
    nu = interp1d(r, nu_r, kind='linear', bounds_error=False, fill_value=0.0)

    # STEP 2:
    sVec = np.arange(0.0, sMax, ds)
    # sVec is the vector of positions along the (half) line of sight chord (km). Multiply by 2 later.
    sVec2 = sVec ** 2  # Get the square for convenience

    # STEP 3: loop over desired range of impact parameters
    r_impVec = np.arange(rSurf, rMax, dr)  # calcualte alpha at all of these impact parameters (km)
    r_imp2 = r_impVec ** 2
    n_imp = r_impVec.shape[0]

    alpha0 = np.zeros(n_imp)

    for i in range(n_imp):
        rs = np.sqrt(r_imp2[i] + sVec2)  # Evaluate refractivities along the line of sight
        alpha0[i] = trapz(nu(rs), sVec) * 2.0  # Integrate the refractivities along the line of sight.

    aFun = interp1d(r_impVec, alpha0, kind='linear', bounds_error=False, fill_value=0.0)

    return (r_impVec, alpha0, aFun)  # Return a paired list of impact parameters (km) and alphas.


# --------------------------------------------

def nu_ref_N2(P, T, lam):
    '''Returns the refractivity of N2 gas (the real index of refraction minus 1) at
    pressure = P (microbar), temperature = T (K) and wavelength = lam (microns).'''

    nu_STP = 29.06e-5 * (1.0 + 7.7e-3 / (lam ** 2))

    # Now calculate the number density (particles per cm^3). We'll ratio this to
    # Loschmidt's number (2.687e19 particles per cm^3 at STP) to get the refractivity
    # at P and T. Use the ideal gas law: PV = NkT, or number density = N/V = P/(kT).
    # Use cgs units: k = 1.380658e-16 erg/K. Remember that one microbar is one dyne/cm^2.

    k = 1.380658e-16
    num_den = P / (k * T)

    LosNum = 2.687e19  # particles per cm^3

    nu_ref = (num_den / LosNum) * nu_STP

    # print("num_den, num_den/LosNum, nu: ", num_den, num_den/LosNum, nu_ref)

    return (nu_ref)


# --------------------------------------------

def atm_ap(npts, FOV, aFun, r_solid, lam):
    '''Generates an aperture of phase delays given a model atmosphere for a circular object.
    The ap screen is npts x npts pixels, corresponding to FOV x FOV km. The argument "aFun" is
    a function that generates alpha as a function of radius, probably created with interp1d.

    INPUTS:
    npts: number of points for the aperture screen. Should be a power of 2 (for FFTs)
    FOV: Physical extent of the aperture field (km)
    aFun: a function that returns alpha (integrated line-of-sight refractivity) as a function of radius
    r_solid: The radius of the solid planet (km)
    lam: the wavelength at which the observations were made (microns)

    OUTPUTS:
    phase_ap: the complex aperture, an npts x npts array.
    '''

    y, x = np.indices((npts, npts)) - npts / 2
    rpix = np.sqrt(y ** 2 + x ** 2)  # radii from the center, in pixels

    lam_km = lam / 10.0 ** 9  # convert microns to km
    # res = np.sqrt(lam_km * D_km/npts)
    # FOV = np.sqrt(npts * lam_km * D_km)
    res = FOV / npts  # sampling resolution, in km per pixel

    print("FOV, resolution (km): ", FOV, res)

    r_km = res * rpix  # radius of every point from the center (km)
    k = 2 * np.pi / lam_km

    phase = np.zeros((npts, npts), dtype="float64")

    atm = np.where(r_km >= r_solid)
    a = aFun(r_km[atm])

    phase[atm] = k * a  # phase delay is alpha times the wavenumber (plus an arbitrary constant)

    phase_ap = np.cos(phase) + 1j * np.sin(phase)

    solid = np.where(r_km < r_solid)
    phase_ap[solid] *= 0.0

    return (phase_ap)

# --------------------------------------------

def atm_ap_2(npts, FOV, aFun, r_solid, lam, e):
    '''Generates an aperture of phase delays given a model atmosphere for a circular object.
    The ap screen is npts x npts pixels, corresponding to FOV x FOV km. The argument "aFun" is
    a function that generates alpha as a function of radius, probably created with interp1d.

    Update in v2: the eccentricity is also added in, which deforms the whole thing by a factor.
    The eccentricity would be the additional parameter specified.
    NOTE: We are keeping that ab=r_solid**2

    INPUTS:
    npts: number of points for the aperture screen. Should be a power of 2 (for FFTs)
    FOV: Physical extent of the aperture field (km)
    aFun: a function that returns alpha (integrated line-of-sight refractivity) as a function of radius
    r_solid: The radius of the solid planet (km)
    lam: the wavelength at which the observations were made (microns)

    OUTPUTS:
    phase_ap: the complex aperture, an npts x npts array.
    '''
    a = r_solid / np.power((1 - e ** 2), 0.25)
    b = np.sqrt(1 - e ** 2) * a



    y, x = np.indices((npts, npts)) - npts / 2
    rpix = np.sqrt(y ** 2 + x ** 2)  # radii from the center, in pixels

    lam_km = lam / 10.0 ** 9  # convert microns to km
    # res = np.sqrt(lam_km * D_km/npts)
    # FOV = np.sqrt(npts * lam_km * D_km)
    res = FOV / npts  # sampling resolution, in km per pixel

    print("FOV, resolution (km): ", FOV, res)

    r_km = res * rpix  # radius of every point from the center (km)
    k = 2 * np.pi / lam_km

    phase = np.zeros((npts, npts), dtype="float64")
    stretch = np.zeros((npts, npts), dtype="float64") # the stretch factor
    mid = int(npts / 2)
    for i in range(mid):
        for j in range(mid):
            tmp = np.arctan((i + 0.5) / (j + 0.5))
            c = np.cos(tmp)
            s = np.sin(tmp)
            stretch[i + mid, j + mid] = np.sqrt((a * c) ** 2 + (b * s) ** 2) / r_solid
            stretch[mid - 1 - i, mid - 1 - j] = stretch[i + mid, j + mid]
            stretch[i + mid, mid - 1 - j] = stretch[i + mid, j + mid]
            stretch[mid - 1 - i, j + mid] = stretch[i + mid, j + mid]
    r_km = r_km * stretch
    atm = np.where(r_km >= r_solid)
    solid = np.where(r_km < r_solid)
    a = aFun(r_km[atm])

    phase[atm] = k * a  # phase delay is alpha times the wavenumber (plus an arbitrary constant)

    phase_ap = np.cos(phase) + 1j * np.sin(phase)


    phase_ap[solid] *= 0.0

    return (phase_ap, phase/k)

# --------------------------------------------


def atm_ap_3_1(npts, FOV, aFun, r_solid, lam, e):
    '''Generates an aperture of phase delays given a model atmosphere for a circular object.
    The ap screen is npts x npts pixels, corresponding to FOV x FOV km. The argument "aFun" is
    a function that generates alpha as a function of radius, probably created with interp1d.

    Update in v3.1: while the sphere is kept constant, the atmosphere is changed to be elliptical
    (Here: Strech happens from the center.)

    INPUTS:
    npts: number of points for the aperture screen. Should be a power of 2 (for FFTs)
    FOV: Physical extent of the aperture field (km)
    aFun: a function that returns alpha (integrated line-of-sight refractivity) as a function of radius
    r_solid: The radius of the solid planet (km)
    lam: the wavelength at which the observations were made (microns)

    OUTPUTS:
    phase_ap: the complex aperture, an npts x npts array.
    '''
    a = r_solid / np.power((1 - e ** 2), 0.25)
    b = np.sqrt(1 - e ** 2) * a



    y, x = np.indices((npts, npts)) - npts / 2
    rpix = np.sqrt(y ** 2 + x ** 2)  # radii from the center, in pixels

    lam_km = lam / 10.0 ** 9  # convert microns to km
    # res = np.sqrt(lam_km * D_km/npts)
    # FOV = np.sqrt(npts * lam_km * D_km)
    res = FOV / npts  # sampling resolution, in km per pixel

    print("FOV, resolution (km): ", FOV, res)

    r_km = res * rpix  # radius of every point from the center (km)
    k = 2 * np.pi / lam_km

    phase = np.zeros((npts, npts), dtype="float64")
    stretch = np.zeros((npts, npts), dtype="float64") # the stretch factor
    mid = int(npts / 2)
    for i in range(mid):
        for j in range(mid):
            tmp = np.arctan((i + 0.5) / (j + 0.5))
            c = np.cos(tmp)
            s = np.sin(tmp)
            stretch[i + mid, j + mid] = np.sqrt((a * c) ** 2 + (b * s) ** 2) / r_solid
            stretch[mid - 1 - i, mid - 1 - j] = stretch[i + mid, j + mid]
            stretch[i + mid, mid - 1 - j] = stretch[i + mid, j + mid]
            stretch[mid - 1 - i, j + mid] = stretch[i + mid, j + mid]
    atm = np.where(r_km >= r_solid)
    solid = np.where(r_km < r_solid)
    r_km = r_km * stretch
    a = aFun(r_km[atm])

    phase[atm] = k * a  # phase delay is alpha times the wavenumber (plus an arbitrary constant)

    phase_ap = np.cos(phase) + 1j * np.sin(phase)

    phase_ap[solid] *= 0.0

    return (phase_ap)

# --------------------------------------------

def atm_ap_3_2(npts, FOV, aFun, r_solid, lam, e):
    '''Generates an aperture of phase delays given a model atmosphere for a circular object.
    The ap screen is npts x npts pixels, corresponding to FOV x FOV km. The argument "aFun" is
    a function that generates alpha as a function of radius, probably created with interp1d.

    Update in v3.2: while the sphere is kept constant, the atmosphere is changed to be elliptical
    (Here: Stretch happens from the surface (i.e. stretched by distance from surface * the scale).)

    INPUTS:
    npts: number of points for the aperture screen. Should be a power of 2 (for FFTs)
    FOV: Physical extent of the aperture field (km)
    aFun: a function that returns alpha (integrated line-of-sight refractivity) as a function of radius
    r_solid: The radius of the solid planet (km)
    lam: the wavelength at which the observations were made (microns)

    OUTPUTS:
    phase_ap: the complex aperture, an npts x npts array.
    '''
    a = r_solid / np.power((1 - e ** 2), 0.25)
    b = np.sqrt(1 - e ** 2) * a



    y, x = np.indices((npts, npts)) - npts / 2
    rpix = np.sqrt(y ** 2 + x ** 2)  # radii from the center, in pixels

    lam_km = lam / 10.0 ** 9  # convert microns to km
    # res = np.sqrt(lam_km * D_km/npts)
    # FOV = np.sqrt(npts * lam_km * D_km)
    res = FOV / npts  # sampling resolution, in km per pixel

    print("FOV, resolution (km): ", FOV, res)

    r_km = res * rpix  # radius of every point from the center (km)
    k = 2 * np.pi / lam_km

    phase = np.zeros((npts, npts), dtype="float64")
    stretch = np.zeros((npts, npts), dtype="float64") # the stretch factor
    mid = int(npts / 2)
    for i in range(mid):
        for j in range(mid):
            tmp = np.arctan((i + 0.5) / (j + 0.5))
            c = np.cos(tmp)
            s = np.sin(tmp)
            stretch[i + mid, j + mid] = np.sqrt((a * c) ** 2 + (b * s) ** 2) / r_solid
            if r_km[i + mid, j + mid]<=r_solid:
                stretch[i + mid, j + mid] = 1 # No stretching inside the circle
            else:
                r_new = (r_km[i + mid, j + mid] - r_solid) * stretch[i + mid, j + mid] + r_solid
                stretch[i + mid, j + mid] = r_new / r_km[i + mid, j + mid] # Record the factor
            stretch[mid - 1 - i, mid - 1 - j] = stretch[i + mid, j + mid]
            stretch[i + mid, mid - 1 - j] = stretch[i + mid, j + mid]
            stretch[mid - 1 - i, j + mid] = stretch[i + mid, j + mid]
    atm = np.where(r_km >= r_solid)
    solid = np.where(r_km < r_solid)
    r_km = r_km * stretch
    a = aFun(r_km[atm])

    phase[atm] = k * a  # phase delay is alpha times the wavenumber (plus an arbitrary constant)

    phase_ap = np.cos(phase) + 1j * np.sin(phase)

    phase_ap[solid] *= 0.0

    return (phase_ap)

# --------------------------------------------

# --------------------------------------------

# --------------------------------------------

def occ_lc(ap):
    '''Calculates the E-field in the observers plane for an occultation by an object with a complex aperture (ap).'''

    # Basic Idea:
    # This routine uses Trester's modified aperture function to calculate the E-field
    # in the far field. The modified aperture function is simply the complex aperture
    # times exp((ik/(2z0)) * (x0**2 + y0**2)). Since k = 2pi/lam and npts = lam*z0, this
    # exponential term is equivalent to exp((i*pi/npts) * (x0**2 + y0**2)).

    npts = ap.shape[0]  # Assume a square aperture
    n2 = int(npts / 2)

    (y, x) = np.indices((npts, npts)) - 1.0 * n2
    eTerm = np.exp(((np.pi * 1j) / npts) * ((y) ** 2 + (x) ** 2))
    M = ap * eTerm

    E_obs = np.fft.fft2(M)
    pObs = np.real(E_obs * np.conj(E_obs))

    return (np.roll(np.roll(pObs, n2, 0), n2, 1))